packages feed

prob-fx 0.1.0.1 → 0.1.0.2

raw patch · 30 files changed

+1268/−987 lines, 30 filesnew-component:exe:prob-fxPVP: minor bump suggested

API additions: PVP suggests at least a minor version bump

API changes (from Hackage documentation)

+ Effects.Dist: pattern Obs :: Member Observe es => PrimDist x -> x -> Addr -> EffectSum es x

Files

README.md view
@@ -1,49 +1,56 @@ ## ProbFX  #### Prelude-ProbFX is a library for probabilistic programming using algebraic effects that implements the paper [**Modular Probabilistic Models via Algebraic Effects**](https://github.com/min-nguyen/prob-fx/blob/master/paper.pdf) -- this paper provides a comprehensive motivation and walkthrough of this library. To have a more interactive and visual play-around with ProbFX, please see https://github.com/min-nguyen/prob-fx: this corresponds parts of the paper to the implementation, and also provides an executable version of ProbFX as a script!+ProbFX is a library for probabilistic programming using algebraic effects that implements the paper [**Modular Probabilistic Models via Algebraic Effects**](https://github.com/min-nguyen/prob-fx/blob/main/paper.pdf) -- this paper provides a comprehensive motivation and walkthrough of this library. To have a more interactive and visual play-around with ProbFX, please see the [**artifact**](https://github.com/min-nguyen/prob-fx/tree/artifact) branch: this corresponds parts of the paper to the implementation, and also provides an executable version of ProbFX as a script.  #### Description-ProbFx is a PPL that places emphasis on being able to define modular and reusable probabilistic models, where the decision to `sample` or `observe` against a random variable or distribution of a model is delayed until the point of execution; this allows a model to be defined just *once* and then reused for a variety of applications. We also implement a compositional approach towards model execution (inference) by using effect handlers. +ProbFx is a PPL that places emphasis on being able to define modular and reusable probabilistic models, where the decision to `sample` or `observe` against a random variable or distribution of a model is delayed until the point of execution; this allows a model to be defined just *once* and then reused for a variety of applications. We also implement a compositional approach towards model execution (inference) by using effect handlers.  #### Building and executing models -A large number of example ProbFX programs are documented [here](https://github.com/min-nguyen/prob-fx/tree/hackage/examples), showing how to define and then execute a probabilistic model. +A large number of example ProbFX programs are documented in the [**examples**](https://github.com/min-nguyen/prob-fx/tree/main/examples) directory, showing how to define and then execute a probabilistic model.  In general, the process is:  1. Define an appropriate model of type `Model env es a`, and (optionally) a corresponding model environment type `env`. -    For example, a logistic regression model that takes a list of `Double`s as inputs and generates a list of `Bool`s:-    ```haskell +    For example, a logistic regression model that takes a list of `Double`s as inputs and generates a list of `Bool`s, modelling the probability of an event occurring or not:+    ```haskell+    -- | The model environment type, for readability purposes     type LogRegrEnv =       '[  "y" ':= Bool,   -- ^ output           "m" ':= Double, -- ^ mean           "b" ':= Double  -- ^ intercept       ] -    logRegr +    -- | Logistic regression model+    logRegr       :: (Observable env "y" Bool        , Observables env '["m", "b"] Double)-      => [Double]           -      -> Model env rs [Bool]  +      => [Double]+      -> Model env rs [Bool]     logRegr xs = do-      -- | Specify distribution of model parameters -      m     <- normal 0 5 #m   -      b     <- normal 0 1 #b     -      sigma <- gamma' 1 1      -      -- | Specify distribution of model output +      -- | Specify the distributions of the model parameters+      -- mean+      m     <- normal 0 5 #m+      -- intercept+      b     <- normal 0 1 #b+      -- noise+      sigma <- gamma' 1 1+      -- | Specify distribution of model outputs       let sigmoid x = 1.0 / (1.0 + exp((-1.0) * x))       ys    <- foldM (\ys x -> do+                        -- probability of event occurring                         p <- normal' (m * x + b) sigma+                        -- generate as output whether the event occurs                         y <- bernoulli (sigmoid p) #y-                        return (y:ys)) [] xs-      return (reverse ys)+                        return (ys ++ [y])) [] xs+      return ys     ```-    The `Observables` constraint says that, for example, `"m"` and `"b"` are observable variables in the model environment `env` that may later be provided a trace of observed values of type `Double`. -    -    Calling a primitive distribution such as `normal 0 5 #m` lets us later provide observed values for "m" when executing the model. -    +    The `Observables` constraint says that, for example, `"m"` and `"b"` are observable variables in the model environment `env` that may later be provided a trace of observed values of type `Double`.++    Calling a primitive distribution such as `normal 0 5 #m` lets us later provide observed values for "m" when executing the model.+     Calling a primed variant of primitive distribution such as `gamma' 1 1` will disable observed values from being provided to that distribution.  2. Execute a model under a model environment, using one of the `Inference` library functions.@@ -54,21 +61,21 @@     simulateLogRegr = do       -- | Specify the model inputs       let xs  = map (/50) [(-50) .. 50]-      -- | Specify the model environment +      -- | Specify the model environment           env = (#y := []) <:> (#m := [2]) <:> (#b := [-0.15]) <:> nil       -- | Simulate from logistic regression       (ys, envs) <- SIM.simulate logRegr env xs       return (zip xs ys)     ```-    +     Below performs Metropolis-Hastings inference on the same model, by providing values for the model output `y` and hence *observing* (conditioning against) them, but providing none for the model parameters `m` and `b` and hence *sampling* them.     ```haskell-    -- | Metropolis-Hastings inference +    -- | Metropolis-Hastings inference     inferMHLogRegr :: Sampler [(Double, Double)]     inferMHLogRegr = do       -- | Simulate data from log regression       (xs, ys) <- unzip <$> simulateLogRegr-      -- | Specify the model environment +      -- | Specify the model environment       let env = (#y := ys) <:> (#m := []) <:> (#b := []) <:> nil       -- | Run MH inference for 20000 iterations       mhTrace :: [Env LogRegrEnv] <- MH.mh 20000 logRegr (xs, env) ["m", "b"]@@ -77,15 +84,15 @@           b_samples = concatMap (get #b) mhTrace       return (zip m_samples b_samples)     ```-    One may have noticed by now that *lists* of values are always provided to observable variables in a model environment; each run-time occurrence of that variable will then result in the head value being observed and consumed, and running out of values will default to sampling. +    One may have noticed by now that *lists* of values are always provided to observable variables in a model environment; each run-time occurrence of that variable will then result in the head value being observed and consumed, and running out of values will default to sampling.      Running the function `mh` returns a trace of output model environments, from which we can retrieve the trace of sampled model parameters via `get #m` and `get #b`. These represent the posterior distribution over `m` and `b`. (The argument `["m", "b"]` to `mh` is optional for indicating interest in learning `#m` and `#b` in particular).  3. `Sampler` computations can be evaluated with `sampleIO :: Sampler a -> IO a` to produce an `IO` computation.      ```haskell-    sampleIO simulateLogRegr :: [(Double, Bool)] +    sampleIO simulateLogRegr :: [(Double, Bool)]     ``` -    +
examples/CoinFlip.hs view
@@ -4,24 +4,35 @@ {-# LANGUAGE OverloadedLabels #-} {-# OPTIONS_GHC -Wno-unrecognised-pragmas #-} {-# HLINT ignore "Redundant return" #-}++{- | A coin-flip model for demonstrating how primitive distributions work in ProbFX.+-}+ module CoinFlip where -import Prog-import Effects.ObsReader-import Model-import PrimDist-import Effects.Dist+import Prog ( call )+import Effects.ObsReader ( ObsReader(Ask) )+import Model ( Model(Model), bernoulli, uniform )+import PrimDist ( PrimDist(BernoulliDist, UniformDist) )+import Effects.Dist ( Dist(Dist) ) import Data.Kind (Constraint)-import Env+import Env ( Observables ) --- ** Coin flip model (Section 5) -coinFlip :: (Observables env '["p"] Double, Observables env '[ "y"] Bool) => Model env es Bool+{- | A coin-flip model that draws a coin-bias @p@ between 0 and 1 from a uniform+     distribution, and uses this to draw a boolean @y@ representing heads or tails.+-}+coinFlip+  :: (Observables env '["p"] Double+    , Observables env '[ "y"] Bool)+  => Model env es Bool coinFlip = do   p <- uniform 0 1 #p   y <- bernoulli p #y   return y --- Desugared coin flip model+{- | A desugared version of the above coin-flip model, after inlining the functions+     @uniform@ and @bernoulli@.+-} coinFlip' :: forall env es. (Observables env '["p"] Double, Observables env '[ "y"] Bool) => Model env es Bool coinFlip' = Model $ do   maybe_p  <- call (Ask @env #p)
examples/HMM.hs view
@@ -7,50 +7,71 @@ {-# OPTIONS_GHC -Wno-unrecognised-pragmas #-} {-# HLINT ignore "Redundant return" #-} +{- | A variety of possible implementations of a [Hidden Markov Model (HMM)](https://en.wikipedia.org/wiki/Hidden_Markov_model).+-}+ module HMM where -import Model-import Inference.SIM as SIM-import Inference.LW as LW-import Sampler-import Control.Monad+import Model ( Model, bernoulli', binomial, uniform )+import Inference.SIM as SIM ( simulate )+import Inference.LW as LW ( lw )+import Sampler ( Sampler )+import Control.Monad ( (>=>) ) import Data.Kind (Constraint)-import Env-import Util+import Env ( Observables, Observable, Assign((:=)), Env, nil, (<:>) )+import Util ( boolToInt ) --- ** HMM Loop +-- | A HMM environment type HMMEnv =-  '[ "trans_p" ':= Double,-     "obs_p"   ':= Double,-     "y"       ':= Int+  '[ "trans_p" ':= Double,    -- ^ parameter for transitioning between latent states+     "obs_p"   ':= Double,    -- ^ parameter for projecting latent state to observation+     "y"       ':= Int        -- ^ observation    ] -hmmFor :: (Observable env "y" Int, Observables env '["obs_p", "trans_p"] Double) =>-  Int -> Int -> Model env ts Int+{- | HMM as a loop+-}+hmmFor :: (Observable env "y" Int, Observables env '["obs_p", "trans_p"] Double)+  -- | number of HMM nodes+  => Int+  -- | initial HMM latent state+  -> Int+  -- | final HMM latent state+  -> Model env ts Int hmmFor n x = do+  -- Draw transition and observation parameters from prior distributions   trans_p <- uniform 0 1 #trans_p   obs_p   <- uniform 0 1 #obs_p+  -- Iterate over @n@ HMM nodes   let hmmLoop i x_prev | i < n = do+                            -- transition to next latent state                             dX <- boolToInt <$> bernoulli' trans_p                             let x = x_prev + dX+                            -- project latent state to observation                             binomial x obs_p #y                             hmmLoop (i - 1) x                        | otherwise = return x_prev   hmmLoop 0 x +-- | Simulate from a HMM simulateHMM :: Sampler (Int, Env HMMEnv) simulateHMM = do+  -- Specify model inputs   let x_0 = 0; n = 10+  -- Specify model environment       env = #trans_p := [0.5] <:> #obs_p := [0.8] <:> #y := [] <:> nil   SIM.simulate (hmmFor n) env 0 +-- | Perform likelihood-weighting inference over HMM inferLwHMM :: Sampler  [(Env HMMEnv, Double)] inferLwHMM   = do+  -- Specify model inputs   let x_0 = 0; n = 10+  -- Specify model environment       env = #trans_p := [] <:> #obs_p := [] <:> #y := [0, 1, 1, 3, 4, 5, 5, 5, 6, 5] <:> nil   LW.lw 100 (hmmFor n) (x_0, env) --- ** Modular HMM  +{- | A modular HMM.+-} transModel ::  Double -> Int -> Model env ts Int transModel transition_p x_prev = do   dX <- boolToInt <$> bernoulli' transition_p@@ -65,7 +86,7 @@ hmmNode :: (Observable env "y" Int)   => Double -> Double -> Int -> Model env ts Int hmmNode transition_p observation_p x_prev = do-  x_i <- transModel  transition_p x_prev+  x_i <- transModel transition_p x_prev   y_i <- obsModel observation_p x_i   return x_i @@ -76,7 +97,8 @@   obs_p   <- uniform 0 1 #obs_p   foldr (>=>) return (replicate n (hmmNode trans_p obs_p)) x --- ** Higher-order, generic HMM +{- | A higher-order, generic HMM.+-} type TransModel env ts params lat   = params -> lat -> Model env ts lat type ObsModel env ts params lat obs = params -> lat -> Model env ts obs
examples/LDA.hs view
@@ -15,45 +15,76 @@ {-# OPTIONS_GHC -Wno-unrecognised-pragmas #-} {-# HLINT ignore "Use camelCase" #-} +{- | A [Latent Dirichlet Allocation (LDA)](https://en.wikipedia.org/wiki/Latent_Dirichlet_allocation) model+     (or topic model) for learning the distribution over words and topics in a text document.+-}+ module LDA where -import Model-import Sampler-import Control.Monad+import Model ( Model, dirichlet, discrete', categorical )+import Sampler ( Sampler )+import Control.Monad ( replicateM ) import Data.Kind (Constraint)-import Env+import Env ( Observables, Observable(get), Assign((:=)), nil, (<:>) ) import Inference.SIM as SIM ( simulate )-import Inference.MH as MH+import Inference.MH as MH ( mh ) +{- | An LDA environment. --- ** Latent dirichlet allocation (topic model) +     Assuming 1 document with 2 topics and a vocabulary of 4 words,+     the parameters of the model environment would have the following shape:++      θ would be [[prob_topic_1, prob_topic_2]                         -- probabilities of topics in document 1+                 ]++      φ would be [[prob_word_1, prob_word_2, prob_word_3, prob_word_4] -- probabilities of words in topic 1+                  [prob_word_1, prob_word_2, prob_word_3, prob_word_4] -- probabilities of words in topic 2+                 ]+-} type TopicEnv =-  '[ "θ" ':= [Double],-     "φ" ':= [Double],-     "w" ':= String+  '[ "θ" ':= [Double],  -- ^ probabilities of each topic in a document+     "φ" ':= [Double],  -- ^ probabilities of each word in a topic+     "w" ':= String     -- ^ word    ] +-- | Prior distribution for topics in a document+docTopicPrior :: Observable env "θ" [Double]+  -- | number of topics+  => Int+  -- | probability of each topic+  -> Model env ts [Double]+docTopicPrior n_topics = dirichlet (replicate n_topics 1) #θ++-- | Prior distribution for words in a topic topicWordPrior :: Observable env "φ" [Double]-  => [String] -> Model env ts [Double]+  -- | vocabulary+  => [String]+  -- | probability of each word+  -> Model env ts [Double] topicWordPrior vocab   = dirichlet (replicate (length vocab) 1) #φ -docTopicPrior :: Observable env "θ" [Double]-  => Int -> Model env ts [Double]-docTopicPrior n_topics = dirichlet (replicate n_topics 1) #θ---- | Distribution over likely words-wordDist :: Observable env "w" String =>-  [String] -> [Double] -> Model env ts String+-- | A distribution generating words according to their probabilities+wordDist :: Observable env "w" String+  -- | vocabulary+  => [String]+  -- | probability of each word+  -> [Double]+  -- | generated word+  -> Model env ts String wordDist vocab ps =   categorical (zip vocab ps) #w  -- | Distribution over the topics in a document, over the distribution of words in a topic topicModel :: (Observables env '["φ", "θ"] [Double],-                 Observable env "w" String)+               Observable env "w" String)+  -- | vocabulary   => [String]+  -- | number of topics   -> Int+  -- | number of words   -> Int+  -- | generated words   -> Model env ts [String] topicModel vocab n_topics n_words = do   -- Generate distribution over words for each topic@@ -67,41 +98,52 @@ -- | Topic distribution over many topics topicModels :: (Observables env '["φ", "θ"] [Double],                Observable env "w" String)-  => [String]  -- Possible vocabulary in a document-  -> Int       -- Assumed number of topics in a document-  -> [Int]     -- Number of words in a document+  -- | vocabulary+  => [String]+  -- | number of topics+  -> Int+  -- | number of words for each document+  -> [Int]+  -- | generated words for each document   -> Model env ts [[String]] topicModels vocab n_topics doc_words = do   mapM (topicModel vocab n_topics) doc_words --- | Simulating from topic model--- | Possible vocabulary++-- | Example possible vocabulary vocab :: [String] vocab = ["DNA", "evolution", "parsing", "phonology"] +-- | Simulating from topic model simLDA :: Sampler [String] simLDA = do-  let n_words = 100+  -- Specify model inputs+  let n_words  = 100+      n_topics = 2+  -- Specify model environment       env_in = #θ := [[0.5, 0.5]] <:>                #φ := [[0.12491280814569208,1.9941599739151505e-2,0.5385152817942926,0.3166303103208638],                       [1.72605174564027e-2,2.9475900240868515e-2,9.906011619752661e-2,0.8542034661052021]] <:>                #w := [] <:> nil-  (words, env_out) <- SIM.simulate (topicModel vocab 2) env_in n_words+  -- Simulate from topic model+  (words, env_out) <- SIM.simulate (topicModel vocab n_topics) env_in n_words   return words --- | MH Inference from topic model--- | Document of words to perform inference over+-- | Example document of words topic_data :: [String] topic_data     = ["DNA","evolution","DNA","evolution","DNA","evolution","DNA","evolution","DNA","evolution", "parsing", "phonology", "DNA","evolution", "DNA", "parsing", "evolution","phonology", "evolution", "DNA","DNA","evolution","DNA","evolution","DNA","evolution","DNA","evolution","DNA","evolution", "parsing", "phonology", "DNA","evolution", "DNA", "parsing", "evolution","phonology", "evolution", "DNA","DNA","evolution","DNA","evolution","DNA","evolution","DNA","evolution","DNA","evolution", "parsing", "phonology", "DNA","evolution", "DNA", "parsing", "evolution","phonology", "evolution", "DNA","DNA","evolution","DNA","evolution","DNA","evolution","DNA","evolution","DNA","evolution", "parsing", "phonology", "DNA","evolution", "DNA", "parsing", "evolution","phonology", "evolution", "DNA","DNA","evolution","DNA","evolution","DNA","evolution","DNA","evolution","DNA","evolution", "parsing", "phonology", "DNA","evolution", "DNA", "parsing", "evolution","phonology", "evolution", "DNA"] +-- | MH inference from topic model mhLDA :: Sampler ([[Double]], [[Double]]) mhLDA  = do-  -- Do MH inference over the topic model using the above data+  -- Specify model inputs   let n_words   = 100       n_topics  = 2+  -- Specify model environment       env_mh_in = #θ := [] <:>  #φ := [] <:> #w := topic_data <:> nil+  -- Run MH for 500 iterations   env_mh_outs <- MH.mh 500 (topicModel vocab n_topics) (n_words, env_mh_in) ["φ", "θ"]-  -- Draw the most recent sampled parameters from MH+  -- Draw the most recent sampled parameters from the MH trace   let env_pred   = head env_mh_outs       θs         = get #θ env_pred       φs         = get #φ env_pred
examples/LinRegr.hs view
@@ -1,4 +1,3 @@- {-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MultiParamTypeClasses #-}@@ -7,61 +6,78 @@ {-# OPTIONS_GHC -Wno-unrecognised-pragmas #-} {-# HLINT ignore "Redundant return" #-} {-# LANGUAGE MonoLocalBinds #-}+{-# LANGUAGE TypeOperators #-} +{- | A linear regression model, assuming a linear relationship between x and y co-ordinates.+-}+ module LinRegr where -import Prog-import Effects.ObsReader-import Effects.Writer-import Model-import Inference.SIM as SIM-import Inference.LW as LW-import Inference.MH as MH-import Effects.Dist-import Effects.Lift-import Sampler-import Control.Monad+import Model ( Model, normal, uniform )+import Inference.SIM as SIM ( simulate )+import Inference.LW as LW ( lw )+import Inference.MH as MH ( mh )+import Sampler ( Sampler ) import Data.Kind (Constraint)-import Env-import Util+import Env ( Observables, Observable(get), Assign((:=)), nil, (<:>) ) --- ** (Section 1) Linear regression-linRegr :: forall env rs . Member Sample rs =>-  Observables env '["y", "m", "c", "σ"] Double =>-  Double -> Model env rs Double+-- | Linear regression environment+type LinRegrEnv =+    '[  "m" ':= Double, -- ^ gradient+        "σ" ':= Double, -- ^ noise+        "c" ':= Double, -- ^ intercept+        "y" ':= Bool    -- ^ output+     ]++-- | Linear regression model+linRegr :: Observables env ["y", "m", "c", "σ"] Double+  -- x co-ordinate+  => Double+  -- y co-ordinate+  -> Model env rs Double linRegr x = do+  -- Draw prior   m <- normal 0 3 #m   c <- normal 0 5 #c   σ <- uniform 1 3 #σ   y <- normal (m * x + c) σ #y   return y --- ** (Section 1, Fig 1a) SIM from linear regression+-- | Simulate from linear regression simulateLinRegr :: Sampler [(Double, Double)] simulateLinRegr = do+  -- Specify model inputs   let xs  = [0 .. 100]+  -- Specify model environment       env = (#m := [3.0]) <:> (#c := [0]) <:> (#σ := [1]) <:> (#y := []) <:> nil+  -- Simulate linear regression for each input x   ys_envs <- mapM (SIM.simulate linRegr env) xs   let ys = map fst ys_envs   return (zip xs ys) --- ** (Section 1, Fig 1b) Perform likelihood weighting over linear regression; returns sampled mu values and associated likelihood weightings+-- | Likelihood weighting over linear regression; returns sampled mu values and associated likelihood weightings inferLwLinRegr :: Sampler [(Double, Double)] inferLwLinRegr = do+  -- Specify model inputs   let xs  = [0 .. 100]+  -- Specify model environments and pair with model input       xys = [(x, env) | x <- xs, let env = (#m := []) <:> (#c := []) <:> (#σ := []) <:> (#y := [3*x]) <:> nil]+  -- Run LW for 200 iterations on each pair of model input and environment   lwTrace <- mapM (LW.lw 200 linRegr) xys-  let -- Get output of LW and extract mu samples-      (env_outs, ps) = unzip $ concat lwTrace+   -- Get the sampled values of mu and their likelihood-weighting+  let (env_outs, ps) = unzip $ concat lwTrace       mus = concatMap (get #m) env_outs   return $ zip mus ps --- Perform Metropolis-Hastings inference over linear regression+-- | Perform Metropolis-Hastings inference over linear regression inferMhLinRegr :: Sampler [Double] inferMhLinRegr = do+  -- Specify model inputs   let xs  = [0 .. 100]+  -- Specify model environments and pair with model input       xys = [(x, env) | x <- xs, let env = (#m := []) <:> (#c := []) <:> (#σ := []) <:> (#y := [3*x]) <:> nil]+  -- Run MH for 100 iterations on each pair of model input and environment   mhTrace <- concat <$> mapM (\xy -> MH.mh 100 linRegr xy ["m", "c"]) xys-  let -- Get output of MH and extract mu samples-      mus = concatMap (get #m) mhTrace+  -- Get the sampled values of mu+  let mus = concatMap (get #m) mhTrace   return mus
examples/LogRegr.hs view
@@ -1,12 +1,15 @@ {-# LANGUAGE RankNTypes #-} {-# LANGUAGE DataKinds #-} {-# LANGUAGE TypeOperators #-}-{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE OverloadedLabels #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE MonoLocalBinds #-} +{- | A logistic regression model, modelling the probability of an event occurring or not.+-}+ module LogRegr where  import Control.Monad ( foldM )@@ -17,33 +20,46 @@ import Inference.MH as MH ( mh ) import Inference.LW as LW ( lw ) --- | Logistic regression model+{- | Logistic regression environment.+     This type definition is for readability purposes and is not used anywhere.+-} type LogRegrEnv =     '[  "y" ':= Bool,   -- ^ output         "m" ':= Double, -- ^ mean         "b" ':= Double  -- ^ intercept      ] -sigmoid :: Double -> Double-sigmoid x = 1 / (1 + exp((-1) * x))--logRegr :: forall rs env.+-- | Logistic regression model+logRegr  -- Specify the "observable variables" that may later be provided observed values- (Observable env "y" Bool, Observables env '["m", "b"] Double) => - [Double] -> Model env rs [Bool]+ :: (Observable env "y" Bool, Observables env '["m", "b"] Double)+ -- | Model inputs+ => [Double]+ -- | Event occurrences+ -> Model env rs [Bool] logRegr xs = do   -- Specify model parameter distributions-  m     <- normal 0 5 #m    -- Annotating with the observable variable #m lets us later provide observed values for m-  b     <- normal 0 1 #b     -  sigma <- gamma' 1 1       -- One can use primed variants of distributions to disable later providing observed values to that variable+  {- Annotating with the observable variable #m lets us later provide observed+     values for m. -}+  m     <- normal 0 5 #m+  b     <- normal 0 1 #b+  {- One can use primed variants of distributions which don't require observable+     variables to be provided. This disables being able to later provide+     observed values to that variable. -}+  sigma <- gamma' 1 1   -- Specify model output distributions-  ls    <- foldM (\ls x -> do-                     y <- normal' (m * x + b) sigma-                     l <- bernoulli (sigmoid y) #y-                     return (l:ls)) [] xs-  return (reverse ls)+  ys    <- foldM (\ys x -> do+                    -- probability of event occurring+                    p <- normal' (m * x + b) sigma+                    -- generate as output whether the event occurs+                    y <- bernoulli (sigmoid p) #y+                    return (ys ++ [y])) [] xs+  return ys --- | SIM from logistic regression+sigmoid :: Double -> Double+sigmoid x = 1 / (1 + exp((-1) * x))++-- | Simulate from logistic regression simulateLogRegr :: Sampler [(Double, Bool)] simulateLogRegr = do   -- First declare the model inputs@@ -75,7 +91,9 @@   (xs, ys) <- unzip <$> simulateLogRegr   let -- Define an environment for inference, providing observed values for the model outputs       env = (#y := ys) <:> (#m := []) <:> (#b := []) <:> nil-  -- Run MH inference for 20000 iterations; the ["m", "b"] is optional for indicating interest in learning #m and #b in particular, causing other variables to not be resampled (unless necessary) during MH.+  -- Run MH inference for 20000 iterations+  {- The agument ["m", "b"] is optional for indicating interest in learning #m and #b in particular,+     causing other variables to not be resampled (unless necessary) during MH. -}   mhTrace :: [Env LogRegrEnv] <- MH.mh 50000 logRegr (xs, env) ["m", "b"]   -- Retrieve values sampled for #m and #b during MH   let m_samples = concatMap (get #m) mhTrace
examples/Main.hs view
@@ -3,13 +3,13 @@  module Main where -import LinRegr-import LogRegr-import SIR-import LDA-import Radon-import School-import Sampler+import LinRegr ( simulateLinRegr, inferLwLinRegr, inferMhLinRegr )+import LogRegr ( simulateLogRegr, inferLwLogRegr, inferMHLogRegr )+import SIR ( simulateSIR, inferSIR, simulateSIRS, simulateSIRSV )+import LDA ( simLDA, mhLDA )+import Radon ( simRadon, mhRadon )+import School ( mhSchool )+import Sampler ( sampleIO ) import System.Environment (getArgs)  printThenWrite :: Show a => a -> IO ()
examples/Radon.hs view
@@ -7,22 +7,34 @@ {-# OPTIONS_GHC -Wno-unrecognised-pragmas #-} {-# HLINT ignore "Redundant return" #-} +{- | A [case study](https://docs.pymc.io/en/v3/pymc-examples/examples/case_studies/multilevel_modeling.html)+    by Gelman and Hill as a hierarchical linear regression model, modelling the relationship between radon levels+    in households in different counties and whether these houses contain basements.+-}+ module Radon where -import Control.Monad-import Model-import Env-import Sampler-import DataSets-import Inference.SIM as SIM-import Inference.MH as MH-import Util+import Control.Monad ( replicateM )+import Model ( Model, normal, halfNormal, halfCauchy' )+import Env ( Observables, Observable(get), Assign((:=)), Env(ENil), (<:>) )+import Sampler ( Sampler, liftS )+import DataSets ( n_counties, logRadon, countyIdx, dataFloorValues )+import Inference.SIM as SIM ( simulate )+import Inference.MH as MH ( mh )+import Util ( findIndexes ) --- | Hierarchical Linear Regression [https://docs.pymc.io/en/v3/pymc-examples/examples/case_studies/multilevel_modeling.html]-type HLREnv =-  '[ "mu_a" ':= Double, "mu_b" ':= Double, "sigma_a" ':= Double, "sigma_b" ':= Double,-     "a" ':= Double, "b" ':= Double, "log_radon" ':= Double]+-- | Radon model environment+type RadonEnv =+  '[ "mu_a" ':= Double      -- ^ hyperprior mean for county intercept+   , "mu_b" ':= Double      -- ^ hyperprior mean for county gradient+   , "sigma_a" ':= Double   -- ^ hyperprior noise for county intercept+   , "sigma_b" ':= Double   -- ^ hyperprior noise for county gradient+   , "a" ':= Double         -- ^ intercept for a county+   , "b" ':= Double         -- ^ gradient for a county+   , "log_radon" ':= Double -- ^ log-radon level for a house+   ] +-- | Prior distribution over model hyperparameters radonPrior :: Observables env '["mu_a", "mu_b", "sigma_a", "sigma_b"] Double   => Model env es (Double, Double, Double, Double) radonPrior = do@@ -32,9 +44,17 @@   sigma_b <- halfNormal 5 #sigma_b   return (mu_a, sigma_a, mu_b, sigma_b) --- n counties = 85, len(floor_x) = 919, len(county_idx) = 919+-- | The Radon model+-- | We have predefined parameters: n counties = 85, len(floor_x) = 919, len(county_idx) = 919 radonModel :: Observables env '["mu_a", "mu_b", "sigma_a", "sigma_b", "a", "b", "log_radon"] Double-  => Int -> [Int] -> [Int] -> () -> Model env es [Double]+  -- | number of counties+  => Int+  -- | whether each house has a basement (1) or not (0)+  -> [Int]+  -- | the county (as an integer) each house belongs to+  -> [Int]+  -> ()+  -> Model env es [Double] radonModel n_counties floor_x county_idx _ = do   (mu_a, sigma_a, mu_b, sigma_b) <- radonPrior   -- Intercept for each county@@ -54,36 +74,45 @@   radon_like <- mapM (\rad_est -> normal rad_est eps #log_radon) radon_est   return radon_like -mkRecordHLR :: ([Double], [Double], [Double], [Double], [Double], [Double], [Double]) -> Env HLREnv+mkRecordHLR :: ([Double], [Double], [Double], [Double], [Double], [Double], [Double]) -> Env RadonEnv mkRecordHLR (mua, mub, siga, sigb, a, b, lograds) =    #mu_a := mua <:> #mu_b := mub <:> #sigma_a := siga <:> #sigma_b := sigb <:> #a := a <:> #b := b <:> #log_radon := lograds <:> ENil +-- | Simulate from the Radon model simRadon :: Sampler ([Double], [Double]) simRadon = do+  -- Specify model environment   let env_in = mkRecordHLR ([1.45], [-0.68], [0.3], [0.2], [], [], [])-  (bs, env_out) <- SIM.simulate  (radonModel n_counties dataFloorValues countyIdx) env_in  ()+  -- Simulate model+  (bs, env_out) <- SIM.simulate (radonModel n_counties dataFloorValues countyIdx) env_in ()+  -- Retrieve and return the radon-levels for houses with basements and those without basements   let basementIdxs      = findIndexes dataFloorValues 0       noBasementIdxs    = findIndexes dataFloorValues 1       basementPoints    = map (bs !!) basementIdxs       nobasementPoints  = map (bs !!) noBasementIdxs   return (basementPoints, nobasementPoints) --- Return posterior for intercepts and gradients+-- | Run MH inference and return posterior for hyperparameter means of intercept and gradient mhRadonpost :: Sampler ([Double], [Double]) mhRadonpost = do+  -- Specify model environment   let env_in = mkRecordHLR ([], [], [], [], [], [], logRadon)+  -- Run MH inference for 2000 iterations   env_outs <- MH.mh 2000 (radonModel n_counties dataFloorValues countyIdx) ((), env_in) ["mu_a", "mu_b", "sigma_a", "sigma_b"]+  -- Retrieve sampled values for model hyperparameters mu_a and mu_b   let mu_a   = concatMap (get #mu_a)  env_outs-      mu_b   = concatMap (get #mu_b)  env_outs+      mu_b   = concatMap (get #mu_a)  env_outs   return (mu_a, mu_b) --- Return predictive posterior for intercepts and gradients+-- | Run MH inference and return predictive posterior for intercepts and gradients mhRadon :: Sampler ([Double], [Double]) mhRadon = do+  -- Specify model environment   let env_in = mkRecordHLR ([], [], [], [], [], [], logRadon)+  -- Run MH inference for 1500 iterations   env_outs <- MH.mh 1500 (radonModel n_counties dataFloorValues countyIdx) ((), env_in) ["mu_a", "mu_b", "sigma_a", "sigma_b"]+  -- Retrieve most recently sampled values for each house's intercept and gradient, a and b   let env_pred   = head env_outs       as         = get #a env_pred       bs         = get #b env_pred-  liftS $ print as   return (as, bs)
examples/SIR.hs view
@@ -1,233 +1,234 @@-{-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeFamilies #-}-{-# LANGUAGE TypeOperators, TypeApplications #-} {-# LANGUAGE OverloadedLabels #-}+{-# LANGUAGE PolyKinds #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE TypeFamilyDependencies #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeOperators #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE TemplateHaskell #-} {-# OPTIONS_GHC -Wno-unrecognised-pragmas #-} {-# HLINT ignore "Redundant return" #-} -module SIR where+{- | This demonstrates:+      - The [SIR](https://en.wikipedia.org/wiki/Compartmental_models_in_epidemiology) model for modelling+        the transition between Susceptible (S), Infected (I), and Recovered (R) individuals during an epidemic.+        We model this as a Hidden Markov Model, where the latent states are the true values of S, I, and R,+        and the observations are the reported number of infections (𝜉).+      - Extending the SIR to the SIRS model where recovered individuals (R) can become susceptible (S) again.+      - Extending the SIRS to the SIRSV model where susceptible individuals (S) can become vaccinated (V). -import Prog-import Effects.Writer-import Model-import Inference.SIM as SIM-import Inference.MH as MH-import Sampler-import Env-import Control.Monad+    For convenience, this makes use of the 'Data.Extensible' library for extensible records, and the 'Control.Lens'+    library to record accessors. If the lens notation is unfamiliar, the code below can be corresponded to a less modular+    version the file [SIRNonModular](examples/SIRNonModular.hs).+ -} -import HMM+module SIR where --- ** The SIR model+import Data.Extensible ( mkField, Record, Lookup, type (>:), (@=), (<:), emptyRecord, Assoc ((:>)) )+import Prog ( Member )+import Control.Lens ( (&), (^.), (.~) )+import Effects.Writer ( Writer, tellM, handleWriterM )+import Model ( Model, beta, binomial', gamma, poisson )+import Control.Monad ( (>=>) )+import Env ( Observables, Observable, Assign ((:=)), (<:>), nil, get)+import HMM ( ObsModel, TransModel, hmmGen )+import GHC.TypeLits ( Symbol )+import Data.Kind (Constraint)+import Sampler ( Sampler )+import Inference.SIM as SIM ( simulate )+import Inference.MH as MH ( mh ) -data Popl = Popl {-    s   :: Int, -- ^ Number of people susceptible to infection-    i   :: Int, -- ^ Number of people currently infected-    r   :: Int  -- ^ Number of people recovered from infection-} deriving Show+-- | A type family for conveniently specifying multiple @Record@ fields of the same type+type family Lookups env (ks :: [Symbol]) a :: Constraint where+  Lookups env (x ': xs) a = (Lookup env x a, Lookups env xs a)+  Lookups env '[] a = () +-- | HMM latent states (SIRV)+mkField "s" -- ^ susceptible individuals+mkField "i" -- ^ infected individuals+mkField "r" -- ^ recovered individuals+mkField "v" -- ^ vaccinated individuals++-- | HMM observations (𝜉) i.e. reported infections type Reported = Int --- | SIR transition model-transSI :: TransModel env ts Double Popl-transSI beta (Popl s i r) = do-  let pop = s + i + r-  dN_SI <- binomial' s (1 - exp ((-beta * fromIntegral i) / fromIntegral pop))-  return $ Popl (s - dN_SI) (i + dN_SI) r -transIR :: TransModel env ts Double Popl-transIR gamma (Popl s i r)  = do-  dN_IR <- binomial' i (1 - exp (-gamma))-  return $ Popl s (i - dN_IR) (r + dN_IR)--data TransParamsSIR = TransParamsSIR {-    betaP  :: Double, -- ^ Mean contact rate between susceptible and infected people-    gammaP :: Double  -- ^ Mean recovery rate-}--transSIR :: Member (Writer [Popl]) es -- || Writer effect from Section 5.5-  => TransModel env es TransParamsSIR Popl-transSIR (TransParamsSIR beta gamma) sir = do-  sir' <- (transSI beta >=> transIR gamma) sir-  tellM [sir'] -  return sir'---- | SIR observation model-type ObsParams = Double--obsSIR :: Observable env "𝜉" Int-  => ObsModel env ts Double Popl Reported-obsSIR rho (Popl _ i _)  = do-  i <- poisson (rho * fromIntegral i) #𝜉-  return i+{- | SIR model.+-}  -- | SIR transition prior transPriorSIR :: Observables env '["β",  "γ"] Double-  => Model env ts TransParamsSIR+  => Model env ts (Double, Double) transPriorSIR = do   pBeta  <- gamma 2 1 #β   pGamma <- gamma 1 (1/8) #γ-  return (TransParamsSIR pBeta pGamma)+  return (pBeta, pGamma) +-- | Transition model from S and I+transSI :: Lookups popl '["s", "i", "r"] Int => TransModel env ts Double (Record popl)+transSI  beta popl = do+  let (s_0, i_0, r_0 ) = (popl ^. s,  popl ^. i,  popl ^. r)+      pop = s_0 + i_0 + r_0+  dN_SI <- binomial' s_0 (1 - exp ((-beta * fromIntegral i_0) / fromIntegral pop))+  return $ popl & s .~ (s_0 - dN_SI)+                & i .~ (i_0 + dN_SI)++-- | Transition model from I and R+transIR :: Lookups popl '["i", "r"] Int => TransModel env ts Double (Record popl)+transIR  gamma popl = do+  let (i_0, r_0) = (popl ^. i,  popl ^. r)+  dN_IR <- binomial' i_0 (1 - exp (-gamma))+  return $ popl & i .~ (i_0 - dN_IR)+                & r .~ (r_0 + dN_IR)++-- | Transition model from S to I, and I to R+transSIR :: (Member (Writer [Record popl]) ts, Lookups popl '["s", "i", "r"] Int)+  => TransModel env ts (Double, Double) (Record popl)+transSIR (beta, gamma) popl = do+  popl <- (transSI beta >=> transIR gamma) popl+  tellM [popl]  -- a user effect for writing each latent SIR state to a stream [Record popl]+  return popl+ -- | SIR observation prior obsPriorSIR :: Observables env '["ρ"] Double-  => Model env ts ObsParams-obsPriorSIR = do-  pRho <- beta 2 7 #ρ-  return pRho---- | SIR as HMM-hmmSIR :: (Member (Writer [Popl]) es, Observable env "𝜉" Int, Observables env '["ρ", "β", "γ"] Double)-  => Int -> Popl -> Model env es Popl-hmmSIR = hmmGen transPriorSIR obsPriorSIR transSIR obsSIR+  => Model env ts Double+obsPriorSIR = beta 2 7 #ρ -hmmSIR' :: (Observables env '["𝜉"] Int , Observables env '[ "β" , "γ" , "ρ"] Double) => Int -> Popl -> Model env es (Popl, [Popl])-hmmSIR' n = handleWriterM . hmmSIR n+-- | SIR observation model+obsSIR :: Lookup s "i" Int => Observable env "𝜉" Int+  => ObsModel env ts Double (Record s) Reported+obsSIR rho popl  = do+  let i_0 = popl ^. i+  poisson (rho * fromIntegral i_0) #𝜉 -type SIRenv = '["β" := Double, "γ"  := Double, "ρ"  := Double, "𝜉" := Int]+-- | SIR as HMM+hmmSIR :: forall popl es env.+           (Lookups popl '["s", "i", "r"] Int, Observable env "𝜉" Int, Observables env '["ρ", "β", "γ"] Double)+  => Int -> Record popl -> Model env es (Record popl, [Record popl])+hmmSIR n = handleWriterM . hmmGen transPriorSIR obsPriorSIR transSIR obsSIR n --- ** Simulating from SIR model: ([(s, i, r)], [𝜉])+-- | Simulate from the SIR model simulateSIR :: Sampler ([(Int, Int, Int)], [Reported]) simulateSIR = do-  let sim_env_in = #β := [0.7] <:> #γ := [0.009] <:> #ρ := [0.3] <:> #𝜉 := [] <:> nil-      sir_0      = Popl {s = 762, i = 1, r = 0}-  ((_, sir_trace), sim_env_out) <- SIM.simulate (hmmSIR' 100) sim_env_in sir_0+  -- Specify model input of 762 susceptible and 1 infected+  let sir_0      = #s @= 762 <: #i @= 1 <: #r @= 0 <: emptyRecord+  -- Specify model environment+      sim_env_in = #β := [0.7] <:> #γ := [0.009] <:> #ρ := [0.3] <:> #𝜉 := [] <:> nil+  -- Simulate an epidemic over 100 days+  ((_, sir_trace), sim_env_out) <- SIM.simulate (hmmSIR 100) sim_env_in sir_0+  -- Get the observed infections over 100 days   let 𝜉s :: [Reported] = get #𝜉 sim_env_out-      sirs = map (\(Popl s i recov) -> (s, i, recov)) sir_trace+  -- Get the true SIR values over 100 days+      sirs = map (\sir -> (sir ^. s, sir ^. i, sir ^. r)) sir_trace   return (sirs, 𝜉s) --- ** MH inference from SIR model: ([ρ], [β])+-- | MH inference from SIR model: ([ρ], [β]) inferSIR :: Sampler ([Double], [Double]) inferSIR = do+  -- Simulate some observed infections   𝜉s <- snd <$> simulateSIR-  let mh_env_in = #β := [] <:> #γ := [0.0085] <:> #ρ := [] <:> #𝜉 := 𝜉s <:> nil-      sir_0           = Popl {s = 762, i = 1, r = 0}-  mhTrace <- MH.mh 50000 (hmmSIR' 100) (sir_0, mh_env_in) ["β", "ρ"]+  -- Specify model input of 762 susceptible and 1 infected+  let sir_0     = #s @= 762 <: #i @= 1 <: #r @= 0 <: emptyRecord+  -- Specify model environment+      mh_env_in = #β := [] <:> #γ := [0.0085] <:> #ρ := [] <:> #𝜉 := 𝜉s <:> nil+  -- Run MH inference over 50000 iterations+  mhTrace <- MH.mh 5000 (hmmSIR 100) (sir_0, mh_env_in) ["β", "ρ"]+  -- Get the sampled values for model parameters ρ and β   let ρs = concatMap (get #ρ) mhTrace       βs = concatMap (get #β) mhTrace   return (ρs, βs) ---- ** Modular Extensions to the SIR Model--{- Note that the implementations below aren't as modular as we would like, due to having to redefine the data types Popl and TransParams when adding new variables to the SIR model. The file "SIRModular.hs" shows how one could take steps to resolve this by using extensible records. -}---- || SIRS (resusceptible) model-data TransParamsSIRS = TransParamsSIRS {-    betaP_SIRS  :: Double, -- ^ Mean contact rate between susceptible and infected people-    gammaP_SIRS :: Double, -- ^ Mean recovery rate-    etaP_SIRS   :: Double  -- ^ Rate of resusceptible-}---- | SIRS transition model-transRS :: Double -> Popl -> Model env ts Popl-transRS eta (Popl s i r) = do-  dN_RS <- binomial' r (1 - exp (-eta))-  return $ Popl (s + dN_RS) i (r - dN_RS)--transSIRS :: Member (Writer [Popl]) es-  => TransModel env es TransParamsSIRS Popl-transSIRS (TransParamsSIRS beta gamma eta) sir = do-  sir' <- (transSI beta >=> transIR gamma >=> transRS eta) sir-  tellM [sir']-  return sir'+{- | SIRS model.+-} --- | SIR transition prior+-- | SIRS transition prior transPriorSIRS :: Observables env '["β", "η", "γ"] Double-  => Model env ts TransParamsSIRS+  => Model env ts (Double, Double, Double) transPriorSIRS = do-  TransParamsSIR pBeta pGamma  <- transPriorSIR+  (pBeta, pGamma)  <- transPriorSIR   pEta <- gamma 1 (1/8) #η-  return (TransParamsSIRS pBeta pGamma pEta)+  return (pBeta, pGamma, pEta) +-- | Transition model from S to R+transRS :: Lookups popl '["s", "r"] Int => TransModel env ts Double (Record popl)+transRS eta popl = do+  let (r_0, s_0) = (popl ^. r,  popl ^. s)+  dN_RS <- binomial' r_0 (1 - exp (-eta))+  return $ popl & r .~ (r_0 - dN_RS)+                & s .~ (s_0 + dN_RS)++-- | Transition model from S to I, I to R, and R to S+transSIRS :: Lookups popl '["s", "i", "r"] Int => TransModel env es (Double, Double, Double) (Record popl)+transSIRS (beta, gamma, eta) = transSI beta >=> transIR gamma >=> transRS eta+ -- | SIRS as HMM-hmmSIRS :: (Observables env '["𝜉"] Int, Observables env '["β", "η", "γ", "ρ"] Double) => Int -> Popl -> Model env ts (Popl, [Popl])+hmmSIRS :: (Lookups popl '["s", "i", "r"] Int,+            Observables env '["𝜉"] Int, Observables env '["β", "η", "γ", "ρ"] Double)+  => Int -> Record popl -> Model env es (Record popl, [Record popl]) hmmSIRS n = handleWriterM . hmmGen transPriorSIRS obsPriorSIR transSIRS obsSIR n --- || (Section 3.2, Fig 4b) SIM from SIRS model: ([(s, i, r)], [𝜉])+-- | Simulate from SIRS model: ([(s, i, r)], [𝜉]) simulateSIRS :: Sampler ([(Int, Int, Int)], [Reported]) simulateSIRS = do-  let sim_env_in = #β := [0.7] <:> #γ := [0.009] <:> #η := [0.05] <:> #ρ := [0.3] <:> #𝜉 := [] <:> nil-      sir_0      = Popl {s = 762, i = 1, r = 0}+  -- Specify model input of 762 susceptible and 1 infected+  let sir_0      = #s @= 762 <: #i @= 1 <: #r @= 0 <: emptyRecord+  -- Specify model environment+      sim_env_in = #β := [0.7] <:> #γ := [0.009] <:> #η := [0.05] <:> #ρ := [0.3] <:> #𝜉 := [] <:> nil+  -- Simulate an epidemic over 100 days   ((_, sir_trace), sim_env_out) <- SIM.simulate (hmmSIRS 100) sim_env_in sir_0+  -- Get the observed infections over 100 days   let 𝜉s :: [Reported] = get #𝜉 sim_env_out-      sirs = map (\(Popl s i recov) -> (s, i, recov)) sir_trace+  -- Get the true SIRS values over 100 days+      sirs = map (\sir -> (sir ^. s, sir ^. i, sir ^. r)) sir_trace   return (sirs, 𝜉s)  --- || SIRSV (resusceptible + vacc) model-data TransParamsSIRSV = TransParamsSIRSV {-    betaP_SIRSV  :: Double, -- ^ Mean contact rate between susceptible and infected people-    gammaP_SIRSV :: Double, -- ^ Mean recovery rate-    etaP_SIRSV   :: Double, -- ^ Rate of resusceptible-    omegaP_SIRSV :: Double  -- ^ Vaccination rate-}--data PoplV = PoplV {-    s' :: Int,-    i' :: Int,-    r' :: Int,-    v' :: Int-} deriving Show---- | SIRSV transition models-transSI' :: TransModel env ts Double PoplV-transSI' beta (PoplV s i r v) = do-  let pop = s + i + r + v-  dN_SI <- binomial' s (1 - exp ((-beta * fromIntegral i) / fromIntegral pop))-  return $ PoplV (s - dN_SI) (i + dN_SI) r v--transIR' :: TransModel env ts Double PoplV-transIR' gamma (PoplV s i r v)  = do-  dN_IR <- binomial' i (1 - exp (-gamma))-  return $ PoplV s (i - dN_IR) (r + dN_IR) v--transRS' :: TransModel env es Double PoplV-transRS' eta (PoplV s i r v) = do-  dN_RS <- binomial' r (1 - exp (-eta))-  return $ PoplV (s + dN_RS) i (r - dN_RS) v--transSV' :: TransModel env es Double PoplV-transSV' omega (PoplV s i r v)  = do-  dN_SV <- binomial' s (1 - exp (-omega))-  return $  PoplV (s - dN_SV) i r (v + dN_SV )--transSIRSV :: Member (Writer [PoplV]) ts => TransModel env ts TransParamsSIRSV PoplV-transSIRSV (TransParamsSIRSV beta gamma omega eta) sirv = do-  sirv' <- (transSI' beta  >=>-            transIR' gamma >=>-            transRS' eta   >=>-            transSV' omega) sirv-  tellM [sirv']-  return sirv'+{- | SIRSV model.+-}  -- | SIRSV transition prior transPriorSIRSV :: Observables env '["β", "γ", "ω", "η"] Double-  => Model env ts TransParamsSIRSV+  => Model env ts (Double, Double, Double, Double) transPriorSIRSV  = do-  TransParamsSIRS pBeta pGamma pEta <- transPriorSIRS+  (pBeta, pGamma, pEta) <- transPriorSIRS   pOmega <- gamma 1 (1/16) #ω-  return (TransParamsSIRSV pBeta pGamma pEta pOmega)+  return (pBeta, pGamma, pEta, pOmega) --- | SIRSV observation model-obsSIRSV :: Observable env "𝜉" Int-  => ObsModel env ts Double PoplV Reported-obsSIRSV rho (PoplV _ i _ v)  = do-  i <- poisson (rho * fromIntegral i) #𝜉-  return i+-- | Transition model from S to V+transSV :: Lookups popl '["s", "v"] Int => TransModel env es Double (Record popl)+transSV omega popl  = do+  let (s_0, v_0) = (popl ^. s,  popl ^. v)+  dN_SV <- binomial' s_0 (1 - exp (-omega))+  return $ popl & s .~ (s_0 - dN_SV)+                & v .~ (v_0 + dN_SV) +-- | Transition model from S to I, I to R, R to S, and S to V+transSIRSV :: Lookups popl '["s", "i", "r", "v"] Int => TransModel env ts (Double, Double, Double, Double) (Record popl)+transSIRSV (beta, gamma, eta, omega) =+  transSI beta >=> transIR gamma >=> transRS eta  >=> transSV omega+ -- | SIRSV as HMM-hmmSIRSV ::  (Observables env '["𝜉"] Int, Observables env '["β", "γ", "η", "ω", "ρ"] Double) => Int -> PoplV -> Model env ts (PoplV, [PoplV])-hmmSIRSV n = handleWriterM . hmmGen transPriorSIRSV obsPriorSIR transSIRSV obsSIRSV n+hmmSIRSV :: (Lookups popl '["s", "i", "r", "v"] Int,+             Observables env '["𝜉"] Int, Observables env '["β", "η", "γ", "ω", "ρ"] Double)+  => Int -> Record popl -> Model env es (Record popl, [Record popl])+hmmSIRSV n = handleWriterM . hmmGen transPriorSIRSV obsPriorSIR transSIRSV obsSIR n --- || Simulate from SIRSV model : ([(s, i, r, v)], [𝜉])+-- | Simulate from SIRSV model : ([(s, i, r, v)], [𝜉]) simulateSIRSV :: Sampler ([(Int, Int, Int, Int)], [Reported]) simulateSIRSV = do-  let sim_env_in = #β := [0.7] <:> #γ := [0.009] <:> #η := [0.05] <:> #ω := [0.02] <:> #ρ := [0.3] <:> #𝜉 := [] <:> nil-      sirv_0      = PoplV {s' = 762, i' = 1, r' = 0, v' = 0}-  ((_, sirv_trace), sim_env_out) <- SIM.simulate (hmmSIRSV 100) sim_env_in sirv_0+  -- Specify model input of 762 susceptible and 1 infected+  let sir_0      = #s @= 762 <: #i @= 1 <: #r @= 0 <: #v @= 0 <: emptyRecord+  -- Specify model environment+      sim_env_in = #β := [0.7] <:> #γ := [0.009] <:> #η := [0.05] <:> #ω := [0.02] <:> #ρ := [0.3] <:> #𝜉 := [] <:> nil+  -- Simulate an epidemic over 100 days+  ((_, sir_trace), sim_env_out) <- SIM.simulate (hmmSIRSV 100) sim_env_in sir_0+  -- Get the observed infections over 100 days   let 𝜉s :: [Reported] = get #𝜉 sim_env_out-      sirvs = map (\(PoplV s i recov v) -> (s, i, recov, v)) sirv_trace+  -- Get the true SIRSV values over 100 days+      sirvs = map (\sirv -> (sirv ^. s, sirv ^. i, sirv ^. r, sirv ^. v)) sir_trace   return (sirvs, 𝜉s)
− examples/SIRModular.hs 
@@ -1,152 +0,0 @@-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeFamilies #-}-{-# LANGUAGE TypeOperators, TypeApplications #-}-{-# LANGUAGE OverloadedLabels #-}-{-# LANGUAGE PolyKinds #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE MultiParamTypeClasses #-}-{-# LANGUAGE TypeFamilyDependencies #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE TypeApplications #-}-{-# LANGUAGE TypeOperators #-}-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE FlexibleInstances #-}-{-# LANGUAGE TemplateHaskell #-}-{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}-{-# HLINT ignore "Redundant return" #-}--module SIRModular where--import qualified Data.Extensible as Extensible-import Data.Extensible hiding (Member)-import Prog-import Control.Lens hiding ((:>))-import Effects.Writer-import Model-import Control.Monad-import Env-import HMM-import Data.Extensible (Associated)-import GHC.TypeLits-import Data.Kind (Constraint)--type family Lookups env (ks :: [Symbol]) a :: Constraint where-  Lookups env (x ': xs) a = (Lookup env x a, Lookups env xs a)-  Lookups env '[] a = ()--mkField "s i r v"--type Reported = Int--{- SIR model using extensible records -}---- | SIR transition model-transSI :: Lookups popl '["s", "i", "r"] Int => TransModel env ts Double (Record popl)-transSI  beta popl = do-  let (s_0, i_0, r_0 ) = (popl ^. s,  popl ^. i,  popl ^. r)-      pop = s_0 + i_0 + r_0-  dN_SI <- binomial' s_0 (1 - exp ((-beta * fromIntegral i_0) / fromIntegral pop))-  return $ popl & s .~ (s_0 - dN_SI)-                & i .~ (i_0 + dN_SI)--transIR :: Lookups popl '["i", "r"] Int => TransModel env ts Double (Record popl)-transIR  gamma popl = do-  let (i_0, r_0) = (popl ^. i,  popl ^. r)-  dN_IR <- binomial' i_0 (1 - exp (-gamma))-  return $ popl & i .~ (i_0 - dN_IR)-                & r .~ (r_0 + dN_IR)--transSIR :: (Member (Writer [Record popl]) ts, Lookups popl '["s", "i", "r"] Int)-  => TransModel env ts (Double, Double) (Record popl)-transSIR (beta, gamma) popl = do-  popl <- (transSI beta >=> transIR gamma) popl-  tellM [popl]-  return popl---- | SIR observation model-type ObsParams = Double--obsSIR :: Lookup s "i" Int => Observable env "𝜉" Int-  => ObsModel env ts Double (Record s) Reported-obsSIR rho popl  = do-  let i_0 = popl ^. i-  poisson (rho * fromIntegral i_0) #𝜉---- | SIR transition prior-transPriorSIR :: Observables env '["β",  "γ"] Double-  => Model env ts (Double, Double)-transPriorSIR = do-  pBeta  <- gamma 2 1 #β-  pGamma <- gamma 1 (1/8) #γ-  return (pBeta, pGamma)---- | SIR observation prior-obsPriorSIR :: Observables env '["ρ"] Double-  => Model env ts ObsParams-obsPriorSIR = beta 2 7 #ρ---- | SIR as HMM-hmmSIR :: (Member (Writer [Record popl]) es,-           Lookups popl '["s", "i", "r"] Int, Observable env "𝜉" Int, Observables env '["ρ", "β", "γ"] Double)-  => Int -> Record popl -> Model env es (Record popl, [Record popl])-hmmSIR n  = handleWriterM . hmmGen transPriorSIR obsPriorSIR transSIR obsSIR n--{- SIRS (resusceptible) model -}---- | SIRS transition model-transRS :: Lookups popl '["s", "r"] Int => TransModel env ts Double (Record popl)-transRS eta popl = do-  let (r_0, s_0) = (popl ^. r,  popl ^. s)-  dN_RS <- binomial' r_0 (1 - exp (-eta))-  return $ popl & r .~ (r_0 - dN_RS)-                & s .~ (s_0 + dN_RS)--transSIRS :: Lookups popl '["s", "i", "r"] Int => TransModel env es (Double, Double, Double) (Record popl)-transSIRS (beta, gamma, eta) = transSI beta >=> transIR gamma >=> transRS eta---- | SIRS transition prior-transPriorSIRS :: Observables env '["β", "η", "γ"] Double-  => Model env ts (Double, Double, Double)-transPriorSIRS = do-  (pBeta, pGamma)  <- transPriorSIR-  pEta <- gamma 1 (1/8) #η-  return (pBeta, pGamma, pEta)---- | SIRS as HMM-hmmSIRS :: (Member (Writer [Record popl]) es,-            Lookups popl '["s", "i", "r"] Int,-            Observables env '["𝜉"] Int, Observables env '["β", "η", "γ", "ρ"] Double)-  => Int -> Record popl -> Model env es (Record popl, [Record popl])-hmmSIRS n = handleWriterM . hmmGen transPriorSIRS obsPriorSIR transSIRS obsSIR n--{- SIRSV (resusceptible + vacc) model -}---- | SIRSV transition model-transSV :: Lookups popl '["s", "v"] Int => TransModel env es Double (Record popl)-transSV omega popl  = do-  let (s_0, v_0) = (popl ^. s,  popl ^. v)-  dN_SV <- binomial' s_0 (1 - exp (-omega))-  return $ popl & s .~ (s_0 - dN_SV)-                & v .~ (v_0 + dN_SV)--transSIRSV :: Lookups popl '["s", "i", "r", "v"] Int => TransModel env ts (Double, Double, Double, Double) (Record popl)-transSIRSV (beta, gamma, eta, omega) =-  transSI beta >=> transIR gamma >=> transRS eta  >=> transSV omega---- | SIRSV transition prior-transPriorSIRSV :: Observables env '["β", "γ", "ω", "η"] Double-  => Model env ts (Double, Double, Double, Double)-transPriorSIRSV  = do-  (pBeta, pGamma, pEta) <- transPriorSIRS-  pOmega <- gamma 1 (1/16) #ω-  return (pBeta, pGamma, pEta, pOmega)---- | SIRSV as HMM-hmmSIRSV :: (Member (Writer [Record popl]) es,-             Lookups popl '["s", "i", "r", "v"] Int,-             Observables env '["𝜉"] Int, Observables env '["β", "η", "γ", "ω", "ρ"] Double)-  => Int -> Record popl -> Model env es (Record popl, [Record popl])-hmmSIRSV n = handleWriterM . hmmGen transPriorSIRSV obsPriorSIR transSIRSV obsSIR n
+ examples/SIRNonModular.hs view
@@ -0,0 +1,287 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE TypeOperators, TypeApplications #-}+{-# LANGUAGE OverloadedLabels #-}+{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}+{-# HLINT ignore "Redundant return" #-}++{- |+     This demonstrates:+      - The [SIR](https://en.wikipedia.org/wiki/Compartmental_models_in_epidemiology) model for modelling+        the transition between Susceptible (S), Infected (I), and Recovered (R) individuals during an epidemic.+        We model this as a Hidden Markov Model, where the _latent states_ are the true values of S, I, and R,+        and the _observations_ are the reported number of infections (𝜉).+      - Extending the SIR to the SIRS model where recovered individuals (R) can become susceptible (S) again.+      - Extending the SIRS to the SIRSV model where susceptible individuals (S) can become vaccinated (V).++      Note that the extensions (SIRS and SIRSV) aren't as modular as we would like, due to having to+      redefine the data types Popl and TransParams when adding new variables to the SIR model.+      The file [SIRModular](examples/SIR.hs) shows how one could take steps to resolve this by using+      extensible records.+-}++module SIRNonModular where++import Prog ( Member )+import Effects.Writer ( Writer, tellM, handleWriterM )+import Model ( Model, beta, binomial', gamma, poisson )+import Inference.SIM as SIM ( simulate )+import Inference.MH as MH ( mh )+import Sampler ( Sampler )+import Env ( Observables, Observable(get), Assign((:=)), nil, (<:>) )+import Control.Monad ( (>=>) )+import HMM ( ObsModel, TransModel, hmmGen )++{- | SIR model.+-}++-- | SIR model environment+type SIRenv =+ '[ "β"  := Double  -- ^ mean contact rate between susceptible and infected people+  , "γ"  := Double  -- ^ mean recovery rate+  , "ρ"  := Double  -- ^ mean report rate of infection+  , "𝜉"  := Int     -- ^ number of reported infections+ ]++-- | Latent state+data Popl = Popl {+    s   :: Int, -- ^ number of people susceptible to infection+    i   :: Int, -- ^ number of people currently infected+    r   :: Int  -- ^ number of people recovered from infection+} deriving Show++-- | Transition model parameters+data TransParamsSIR = TransParamsSIR {+    betaP  :: Double, -- ^ mean contact rate between susceptible and infected people+    gammaP :: Double  -- ^ mean recovery rate+}++-- | Observation 𝜉+type Reported = Int++-- | Observation model parameters+type ObsParams = Double++-- | Transition model prior+transPriorSIR :: Observables env '["β",  "γ"] Double+  => Model env ts TransParamsSIR+transPriorSIR = do+  pBeta  <- gamma 2 1 #β+  pGamma <- gamma 1 (1/8) #γ+  return (TransParamsSIR pBeta pGamma)++-- | Transition model between S and I+transSI :: TransModel env ts Double Popl+transSI beta (Popl s i r) = do+  let pop = s + i + r+  dN_SI <- binomial' s (1 - exp ((-beta * fromIntegral i) / fromIntegral pop))+  return $ Popl (s - dN_SI) (i + dN_SI) r++-- | Transition model between I and R+transIR :: TransModel env ts Double Popl+transIR gamma (Popl s i r)  = do+  dN_IR <- binomial' i (1 - exp (-gamma))+  return $ Popl s (i - dN_IR) (r + dN_IR)++-- | Transition model between S, I, and R+transSIR :: Member (Writer [Popl]) es+  => TransModel env es TransParamsSIR Popl+transSIR (TransParamsSIR beta gamma) sir = do+  sir' <- (transSI beta >=> transIR gamma) sir+  tellM [sir']  -- a user effect for writing each latent SIR state to a stream [Popl]+  return sir'++-- | Observation model prior+obsPriorSIR :: Observables env '["ρ"] Double+  => Model env ts ObsParams+obsPriorSIR = do+  pRho <- beta 2 7 #ρ+  return pRho++-- | Observation model from I to 𝜉+obsSIR :: Observable env "𝜉" Int+  => ObsModel env ts Double Popl Reported+obsSIR rho (Popl _ i _)  = do+  i <- poisson (rho * fromIntegral i) #𝜉+  return i++-- | SIR as HMM+hmmSIR :: (Member (Writer [Popl]) es, Observable env "𝜉" Int, Observables env '["ρ", "β", "γ"] Double)+  => Int -> Popl -> Model env es Popl+hmmSIR = hmmGen transPriorSIR obsPriorSIR transSIR obsSIR++-- | Handle the user effect for writing each SIR state to a stream [Popl]+hmmSIR' :: (Observables env '["𝜉"] Int , Observables env '[ "β" , "γ" , "ρ"] Double) => Int -> Popl -> Model env es (Popl, [Popl])+hmmSIR' n = handleWriterM . hmmSIR n++-- | Simulating from SIR model: ([(s, i, r)], [𝜉])+simulateSIR :: Sampler ([(Int, Int, Int)], [Reported])+simulateSIR = do+  -- Specify model input of 762 susceptible and 1 infected+  let sir_0      = Popl {s = 762, i = 1, r = 0}+  -- Specify model environment+      sim_env_in = #β := [0.7] <:> #γ := [0.009] <:> #ρ := [0.3] <:> #𝜉 := [] <:> nil+  -- Simulate an epidemic over 100 days+  ((_, sir_trace), sim_env_out) <- SIM.simulate (hmmSIR' 100) sim_env_in sir_0+  -- Get the observed infections over 100 days+  let 𝜉s :: [Reported] = get #𝜉 sim_env_out+  -- Get the true SIR values over 100 days+      sirs = map (\(Popl s i recov) -> (s, i, recov)) sir_trace+  return (sirs, 𝜉s)++-- | MH inference from SIR model: ([ρ], [β])+inferSIR :: Sampler ([Double], [Double])+inferSIR = do+  -- Simulate some observed infections+  𝜉s <- snd <$> simulateSIR+  -- Specify model input of 762 susceptible and 1 infected+  let sir_0           = Popl {s = 762, i = 1, r = 0}+  -- Specify model environment+      mh_env_in = #β := [] <:> #γ := [0.0085] <:> #ρ := [] <:> #𝜉 := 𝜉s <:> nil+  -- Run MH inference over 50000 iterations+  mhTrace <- MH.mh 5000 (hmmSIR' 100) (sir_0, mh_env_in) ["β", "ρ"]+  -- Get the sampled values for model parameters ρ and β+  let ρs = concatMap (get #ρ) mhTrace+      βs = concatMap (get #β) mhTrace+  return (ρs, βs)+++{- | SIRS model.+-}+-- | Transition model parameters+data TransParamsSIRS = TransParamsSIRS {+    betaP_SIRS  :: Double, -- ^ mean contact rate between susceptible and infected people+    gammaP_SIRS :: Double, -- ^ mean recovery rate+    etaP_SIRS   :: Double  -- ^ rate of resusceptible+}++-- | Transition model prior+transPriorSIRS :: Observables env '["β", "η", "γ"] Double+  => Model env ts TransParamsSIRS+transPriorSIRS = do+  TransParamsSIR pBeta pGamma  <- transPriorSIR+  pEta <- gamma 1 (1/8) #η+  return (TransParamsSIRS pBeta pGamma pEta)++-- | Transition model between R and S+transRS :: Double -> Popl -> Model env ts Popl+transRS eta (Popl s i r) = do+  dN_RS <- binomial' r (1 - exp (-eta))+  return $ Popl (s + dN_RS) i (r - dN_RS)++-- | Transition model between S, to I, to R, and to S+transSIRS :: Member (Writer [Popl]) es+  => TransModel env es TransParamsSIRS Popl+transSIRS (TransParamsSIRS beta gamma eta) sir = do+  sir' <- (transSI beta >=> transIR gamma >=> transRS eta) sir+  tellM [sir']+  return sir'++-- | SIRS as HMM+hmmSIRS :: (Observables env '["𝜉"] Int, Observables env '["β", "η", "γ", "ρ"] Double) => Int -> Popl -> Model env ts (Popl, [Popl])+hmmSIRS n = handleWriterM . hmmGen transPriorSIRS obsPriorSIR transSIRS obsSIR n++-- | Simulate from SIRS model: ([(s, i, r)], [𝜉])+simulateSIRS :: Sampler ([(Int, Int, Int)], [Reported])+simulateSIRS = do+  -- Specify model input of 762 susceptible and 1 infected+  let sir_0      = Popl {s = 762, i = 1, r = 0}+  -- Specify model environment+      sim_env_in = #β := [0.7] <:> #γ := [0.009] <:> #η := [0.05] <:> #ρ := [0.3] <:> #𝜉 := [] <:> nil+  -- Simulate an epidemic over 100 days+  ((_, sir_trace), sim_env_out) <- SIM.simulate (hmmSIRS 100) sim_env_in sir_0+  -- Get the observed infections over 100 days+  let 𝜉s :: [Reported] = get #𝜉 sim_env_out+  -- Get the true SIR values over 100 days+      sirs = map (\(Popl s i recov) -> (s, i, recov)) sir_trace+  return (sirs, 𝜉s)+++{- | SIRSV model.+-}+-- | Transition model parameters+data TransParamsSIRSV = TransParamsSIRSV {+    betaP_SIRSV  :: Double, -- ^ mean contact rate between susceptible and infected people+    gammaP_SIRSV :: Double, -- ^ mean recovery rate+    etaP_SIRSV   :: Double, -- ^ rate of resusceptible+    omegaP_SIRSV :: Double  -- ^ vaccination rate+}++-- | Latent state+data PoplV = PoplV {+    s' :: Int,  -- ^ susceptible individuals+    i' :: Int,  -- ^ infected individuals+    r' :: Int,  -- ^ recovered individuals+    v' :: Int   -- ^ vaccinated individuals+} deriving Show++-- | Transition from S to I+transSI' :: TransModel env ts Double PoplV+transSI' beta (PoplV s i r v) = do+  let pop = s + i + r + v+  dN_SI <- binomial' s (1 - exp ((-beta * fromIntegral i) / fromIntegral pop))+  return $ PoplV (s - dN_SI) (i + dN_SI) r v++-- | Transition from I to R+transIR' :: TransModel env ts Double PoplV+transIR' gamma (PoplV s i r v)  = do+  dN_IR <- binomial' i (1 - exp (-gamma))+  return $ PoplV s (i - dN_IR) (r + dN_IR) v++-- | Transition from R to S+transRS' :: TransModel env es Double PoplV+transRS' eta (PoplV s i r v) = do+  dN_RS <- binomial' r (1 - exp (-eta))+  return $ PoplV (s + dN_RS) i (r - dN_RS) v++-- | Transition from S to V+transSV' :: TransModel env es Double PoplV+transSV' omega (PoplV s i r v)  = do+  dN_SV <- binomial' s (1 - exp (-omega))+  return $  PoplV (s - dN_SV) i r (v + dN_SV )++-- | Transition between S to I, I to R, R to S, and S to V+transSIRSV :: Member (Writer [PoplV]) ts => TransModel env ts TransParamsSIRSV PoplV+transSIRSV (TransParamsSIRSV beta gamma omega eta) sirv = do+  sirv' <- (transSI' beta  >=>+            transIR' gamma >=>+            transRS' eta   >=>+            transSV' omega) sirv+  tellM [sirv']+  return sirv'++-- | Transition model prior+transPriorSIRSV :: Observables env '["β", "γ", "ω", "η"] Double+  => Model env ts TransParamsSIRSV+transPriorSIRSV  = do+  TransParamsSIRS pBeta pGamma pEta <- transPriorSIRS+  pOmega <- gamma 1 (1/16) #ω+  return (TransParamsSIRSV pBeta pGamma pEta pOmega)++-- | Observation model+obsSIRSV :: Observable env "𝜉" Int+  => ObsModel env ts Double PoplV Reported+obsSIRSV rho (PoplV _ i _ v)  = do+  i <- poisson (rho * fromIntegral i) #𝜉+  return i++-- | SIRSV as HMM+hmmSIRSV ::  (Observables env '["𝜉"] Int, Observables env '["β", "γ", "η", "ω", "ρ"] Double) => Int -> PoplV -> Model env ts (PoplV, [PoplV])+hmmSIRSV n = handleWriterM . hmmGen transPriorSIRSV obsPriorSIR transSIRSV obsSIRSV n++-- | Simulate from SIRSV model : ([(s, i, r, v)], [𝜉])+simulateSIRSV :: Sampler ([(Int, Int, Int, Int)], [Reported])+simulateSIRSV = do+  -- Specify model input of 762 susceptible and 1 infected+  let sirv_0      = PoplV {s' = 762, i' = 1, r' = 0, v' = 0}+  -- Specify model environment+      sim_env_in = #β := [0.7] <:> #γ := [0.009] <:> #η := [0.05] <:> #ω := [0.02] <:> #ρ := [0.3] <:> #𝜉 := [] <:> nil+  -- Simulate an epidemic over 100 days+  ((_, sirv_trace), sim_env_out) <- SIM.simulate (hmmSIRSV 100) sim_env_in sirv_0+  -- Get the observed infections over 100 days+  let 𝜉s :: [Reported] = get #𝜉 sim_env_out+  -- Get the true SIRV values over 100 days+      sirvs = map (\(PoplV s i recov v) -> (s, i, recov, v)) sirv_trace+  return (sirvs, 𝜉s)
examples/School.hs view
@@ -6,41 +6,54 @@ {-# LANGUAGE TypeOperators #-} {-# LANGUAGE OverloadedLabels #-} +{- | The Gelman and Hill [8-schools case study](https://cran.r-project.org/web/packages/rstan/vignettes/rstan.html),+     which quantifies the effect of coaching programs from 8 different schools on students' SAT-V scores.+-}+ module School where -import Model-import Inference.MH as MH-import Sampler-import Control.Monad+import Model ( Model, deterministic, normal, normal', halfNormal' )+import Inference.MH as MH ( mh )+import Sampler ( Sampler )+import Control.Monad ( replicateM ) import Data.Kind (Constraint)-import Env+import Env ( Observables, Observable(get), Assign((:=)), Env(ENil), (<:>) ) --- | Hierarchical School Model+-- | School model environment type SchEnv = '[-    "mu"    ':= Double,-    "theta" ':= [Double],-    "y"     ':= Double+    "mu"    ':= Double,   -- ^ effect of general coaching programs on SAT scores+    "theta" ':= [Double], -- ^ variation of each program's effect on SAT scores+    "y"     ':= Double    -- ^ effectiveness on SAT scores   ] +-- | School model schoolModel :: (Observables env '["mu", "y"] Double, Observable env "theta" [Double])-  => Int -> [Double] -> Model env es [Double]+  -- | number of schools+  => Int+  -- | standard errors of each school+  -> [Double]+  -- | effectiveness of each school+  -> Model env es [Double] schoolModel n_schools σs = do   μ   <- normal 0 10 #mu   τ   <- halfNormal' 10   ηs  <- replicateM n_schools (normal' 0 1)   θs  <- deterministic (map ((μ +) . (τ *)) ηs) #theta   ys  <- mapM (\(θ, σ) -> normal θ σ #y) (zip θs σs)-  let h = ""   return θs --- Perform MH inference+-- | Perform MH inference mhSchool :: Sampler ([Double], [[Double]]) mhSchool = do+  -- Specify model inputs   let n_schools = 8       ys        = [28, 8, -3,   7, -1,  1, 18, 12]       sigmas    = [15, 10, 16, 11,  9, 11, 10, 18]+  -- Specify model environment       env       = #mu := [] <:> #theta := [] <:> #y := ys <:> ENil+  -- Run MH inference for 10000 iterations   env_mh_out <- MH.mh 10000 (schoolModel n_schools ) (sigmas, env) ["mu", "theta"]+  -- Retrieve and returns the trace of model parameters mu and theta   let mus    = concatMap (get #mu) env_mh_out       thetas = concatMap (get #theta) env_mh_out   return (mus, thetas)
prob-fx.cabal view
@@ -1,6 +1,6 @@ cabal-version:       3.0 name:                prob-fx-version:             0.1.0.1+version:             0.1.0.2 license:             BSD-3-Clause license-file:        LICENSE.md copyright:           2022 Minh Nguyen@@ -9,10 +9,10 @@ maintainer:          minhnguyen1995@googlemail.com homepage:            https://github.com/min-nguyen/prob-fx/tree/hackage synopsis:            A library for modular probabilistic modelling-description:         +description:   A library for probabilistic programming using algebraic effects. The   emphasis is on modular definitions of probabilistic models, and also-  compositional implementation of model execution (inference) in terms +  compositional implementation of model execution (inference) in terms   of effect handlers.  category:            Statistics@@ -28,7 +28,7 @@   GHC == 9.0.1  library-  exposed-modules:    +  exposed-modules:                       Effects.Dist,                       Effects.Lift,                       Effects.ObsReader,@@ -81,7 +81,7 @@   default-language:    Haskell2010   ghc-options:         -funfolding-use-threshold=16 -fexcess-precision -optc-O3 -optc-ffast-math -executable examples+executable prob-fx   default-language: Haskell2010   hs-source-dirs: examples   main-is:        Main.hs@@ -94,8 +94,8 @@                   Radon,                   School,                   SIR,-                  SIRModular,-  build-depends:  base, +                  SIRNonModular,+  build-depends:  base,                   prob-fx,                   lens,                   extensible
src/Effects/Dist.hs view
@@ -1,21 +1,14 @@-{-# LANGUAGE PatternSynonyms #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE GADTs, TypeOperators #-} {-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE PolyKinds #-}-{-# LANGUAGE TypeFamilies #-}-{-# LANGUAGE ViewPatterns #-}-{-# LANGUAGE FlexibleInstances #-}-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE LambdaCase #-}-{-# LANGUAGE ConstraintKinds #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE TypeOperators #-} -{- | The effect for primitive distributions +{- | The effects for primitive distributions, sampling, and observing. -}  module Effects.Dist (   -- ** Address+  -- $Address     Tag   , Addr   -- ** Dist effect@@ -31,19 +24,19 @@ import Data.Maybe ( fromMaybe ) import Prog ( call, discharge, Member, Prog(..) ) import qualified Data.Map as Map-import qualified Data.Vector as V-import qualified Data.Vector.Unboxed as UV-import qualified OpenSum-import Util ( boolToInt ) import PrimDist ( PrimDist ) --- | An observable variable name assigned to a primitive distribution+{- $Address+   Run-time identifiers for probabilistic operations+-}++-- | An observable variable name assigned to a primitive distribution, representing a compile-time identifier type Tag  = String--- | An observable variable name and the index of its run-time occurrence+-- | An observable variable name and the index of its run-time occurrence, representing a run-time identifier type Addr = (Tag, Int) --- | The Dist effect-data Dist a = Dist +-- | The effect @Dist@ for primitive distributions+data Dist a = Dist   { getPrimDist :: PrimDist a  -- ^ primitive distribution   , getObs :: Maybe a          -- ^ optional observed value   , getTag :: Maybe Tag        -- ^ optional observable variable name@@ -53,24 +46,25 @@   show (Dist d y tag) = "Dist(" ++ show d ++ ", " ++ show y ++ ", " ++ show tag ++ ")"  instance Eq (Dist a) where-  (==) (Dist d1 _ _) (Dist d2 _ _) = d1 == d2 +  (==) (Dist d1 _ _) (Dist d2 _ _) = d1 == d2 --- | An effect for sampling from distirbutions+-- | The effect @Sample@ for sampling from distirbutions data Sample a where-  Sample  :: PrimDist a     -- ^ Distribution to sample from-          -> Addr           -- ^ Address of @Sample@ operation+  Sample  :: PrimDist a     -- ^ distribution to sample from+          -> Addr           -- ^ address of @Sample@ operation           -> Sample a --- | An effect for conditioning against observed values+-- | The effect @Observe@ for conditioning against observed values data Observe a where-  Observe :: PrimDist a     -- ^ Distribution to condition with-          -> a              -- ^ Observed value-          -> Addr           -- ^ Address of @Observe@ operation+  Observe :: PrimDist a     -- ^ distribution to condition with+          -> a              -- ^ observed value+          -> Addr           -- ^ address of @Observe@ operation           -> Observe a --- | Handle the @Dist@ effect to a @Sample@ or @Observe@ effect and assign address+-- | Handle the @Dist@ effect to a @Sample@ or @Observe@ effect and assign an address handleDist :: (Member Sample es, Member Observe es)-        => Prog (Dist : es) a -> Prog es a+  => Prog (Dist : es) a+  -> Prog es a handleDist = loop 0 Map.empty   where   loop :: (Member Sample es, Member Observe es)
src/Effects/Lift.hs view
@@ -1,15 +1,9 @@-{-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GADTs #-}-{-# LANGUAGE MultiParamTypeClasses #-}-{-# LANGUAGE PolyKinds #-} {-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeFamilies #-}-{-# LANGUAGE AllowAmbiguousTypes #-} -{- | For lifting arbitrary monadic computations into an algebraic effect setting+{- | For lifting arbitrary monadic computations into an algebraic effect setting. -}  module Effects.Lift (@@ -20,10 +14,10 @@ import Prog ( call, Member(prj), Prog(..) ) import Data.Function (fix) --- | Lift effect+-- | Lift a monadic computation @m a@ into the effect @Lift m@ newtype Lift m a = Lift (m a) --- | Lift a monadic computation @m a@ into the effect @Lift m@+-- | Wrapper function for calling @Lift@ lift :: (Member (Lift m) es) => m a -> Prog es a lift = call . Lift
src/Effects/ObsReader.hs view
@@ -1,11 +1,13 @@+{-# LANGUAGE AllowAmbiguousTypes #-} {-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeOperators, TypeApplications #-}-{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeOperators #-} -{- | An effect for reading observable variables from a model environment -}+{- | The effect for reading observable variables from a model environment.+-}  module Effects.ObsReader (     ObsReader(..)@@ -18,20 +20,22 @@  -- | The effect for reading observed values from a model environment @env@ data ObsReader env a where-  -- | Given variable @x@ is assigned a list of type @[a]@ in @env@, attempt to retrieve its head value.-  Ask :: Observable env x a       -    => ObsVar x                 -- ^ Variable @x@ to read from-    -> ObsReader env (Maybe a)  -- ^ The head value from @x@'s list+  -- | Given the observable variable @x@ is assigned a list of type @[a]@ in @env@, attempt to retrieve its head value.+  Ask :: Observable env x a+    => ObsVar x                 -- ^ variable @x@ to read from+    -> ObsReader env (Maybe a)  -- ^ the head value from @x@'s list --- | Retrieve the first value from the trace of an observable variable if it exists-ask :: forall env es x a. (Member (ObsReader env) es, Observable env x a) =>  ObsVar x -> Prog es (Maybe a)+-- | Wrapper function for calling @Ask@+ask :: forall env es x a. (Member (ObsReader env) es, Observable env x a)+  => ObsVar x+  -> Prog es (Maybe a) ask x = call (Ask @env x)  -- | Handle the @Ask@ requests of observable variables-handleRead :: -  -- | Initial model environment-     Env env -  -> Prog (ObsReader env ': es) a +handleRead ::+  -- | initial model environment+     Env env+  -> Prog (ObsReader env ': es) a   -> Prog es a handleRead env (Val x) = return x handleRead env (Op op k) = case discharge op of
src/Effects/State.hs view
@@ -4,12 +4,12 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeOperators #-} -{- | The state effect+{- | State effect. -}  module Effects.State (     State(..)-  , get +  , get   , put   , modify   , handleState) where@@ -18,28 +18,29 @@  -- | The state effect data State s a where+  -- | Get the current state   Get :: State s s+  -- | Set the current state   Put :: s -> State s () --- | Get the state-get :: (Member (State s) es) => Prog es s+-- | Wrapper function for @Get@+get :: Member (State s) es => Prog es s get = Op (inj Get) Val --- | Set the state+-- | Wrapper function for @Set@ put :: (Member (State s) es) => s -> Prog es () put s = Op (inj $ Put s) Val --- | Apply a function to the state+-- | Wrapper function for apply a function to the state modify :: Member (State s) es => (s -> s) -> Prog es () modify f = get >>= put . f  -- | Handle the @State s@ effect-handleState :: -  -- | Initial state-     s -  -- | Initial program-  -> Prog (State s ': es) a -  -- | Pure value and final state+handleState+  -- | initial state+  :: s+  -> Prog (State s ': es) a+  -- | (output, final state)   -> Prog es (a, s) handleState s m = loop s m where   loop :: s -> Prog (State s ': es) a -> Prog es (a, s)
src/Effects/Writer.hs view
@@ -4,7 +4,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeOperators #-} -{- | Writer effect+{- | Writer effect. -}  module Effects.Writer (@@ -20,27 +20,32 @@ -- | Writer effect for writing to a strean @w@ data Writer w a where   -- | Write to a stream @w@-  Tell :: -       w  -- ^ Value to write-    -> Writer w ()+  Tell :: w             -- ^ value to write+       -> Writer w () --- | Write to a stream @w@-tell :: Member (Writer w) ts => w -> Prog ts ()+-- | Wrapper for @Tell@+tell :: Member (Writer w) es => w -> Prog es () tell w = Op (inj $ Tell w) Val-    --- | Write to a stream @w@ inside a @Model@++-- | Wrapper for @Tell@ inside @Model@ tellM :: Member (Writer w) es => w -> Model env es () tellM w = Model $ tell w --- | Handle the @Writer@ effect-handleWriter :: forall w ts a . Monoid w => Prog (Writer w ': ts) a -> Prog ts (a, w)+-- | Handle the @Writer@ effect for a stream @w@+handleWriter :: forall w es a. Monoid w+  => Prog (Writer w ': es) a+  -- | (output, final stream)+  -> Prog es (a, w) handleWriter = loop mempty where-  loop ::  w -> Prog (Writer w ': ts) a -> Prog ts (a, w)+  loop ::  w -> Prog (Writer w ': es) a -> Prog es (a, w)   loop w (Val x) = return (x, w)   loop w (Op u k) = case discharge u of     Right (Tell w') -> loop (w `mappend` w') (k ())     Left u'         -> Op u' (loop w . k)  -- | Handle the @Writer@ effect inside a @Model@-handleWriterM :: Monoid w => Model env (Writer w : es) v -> Model env es (v, w)+handleWriterM :: Monoid w+  => Model env (Writer w : es) a+  -- | (output, final stream)+  -> Model env es (a, w) handleWriterM m = Model $ handleWriter $ runModel m
src/Env.hs view
@@ -1,20 +1,21 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE PolyKinds #-} {-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE MultiParamTypeClasses #-}-{-# LANGUAGE TypeFamilyDependencies #-}-{-# LANGUAGE RankNTypes #-} {-# LANGUAGE TypeApplications #-}+{-# LANGUAGE TypeFamilyDependencies #-} {-# LANGUAGE TypeOperators #-}-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE FlexibleInstances #-}-{-# LANGUAGE FlexibleContexts #-} {-# OPTIONS_GHC -Wno-incomplete-patterns #-} -{- | This implements the model environments that users must provide upon running a model; such environments assign traces of values to the "observable variables" (random variables which can be conditioned against) of a model.+{- | This implements the model environments that users must provide upon running a model;+     such environments assign traces of values to the "observable variables" (random+     variables which can be conditioned against) of a model. -} -module Env +module Env   ( -- * Observable variable     ObsVar(..)   , varToStr@@ -33,27 +34,26 @@ import FindElem ( FindElem(..), Idx(..) ) import GHC.OverloadedLabels ( IsLabel(..) ) import GHC.TypeLits ( KnownSymbol, Symbol, symbolVal )-import qualified Data.Vector as V-import qualified GHC.TypeLits as TL import Unsafe.Coerce ( unsafeCoerce ) ----- | Containers for observable variables +-- | Containers for observable variables data ObsVar (x :: Symbol) where   ObsVar :: KnownSymbol x => ObsVar x  -- | Allows the syntax @#x@ to be automatically lifted to the type @ObsVar "x"@. instance (KnownSymbol x, x ~ x') => IsLabel x (ObsVar x') where   fromLabel = ObsVar-  + -- | Convert an observable variable from a type-level string to a value-level string varToStr :: forall x. ObsVar x -> String varToStr ObsVar = symbolVal (Proxy @x) --- * Model Environments +-- * Model Environments --- | A model environment assigning traces (lists) of observed values to observable variables i.e. the type @Env ((x := a) : env)@ indicates @x@ is assigned a value of type @[a]@+{- | A model environment assigning traces (lists) of observed values to observable+     variables i.e. the type @Env ((x := a) : env)@ indicates @x@ is assigned a value+     of type @[a]@.+-} data Env (env :: [Assign Symbol *]) where   ENil  :: Env '[]   ECons :: [a] -> Env env -> Env (x := a : env)
src/FindElem.hs view
@@ -1,13 +1,12 @@-{-# LANGUAGE MultiParamTypeClasses #-}-{-# LANGUAGE TypeFamilyDependencies #-}-{-# LANGUAGE TypeOperators #-}-{-# LANGUAGE PolyKinds #-} {-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleInstances #-}-{-# LANGUAGE UndecidableInstances #-}+{-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE PolyKinds #-} {-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeOperators #-}+{-# LANGUAGE UndecidableInstances #-} -{- | Auxiliary definitions for finding a type in a type-level list. +{- | Auxiliary definitions for finding a type in a type-level list. -}  module FindElem (
src/Inference/LW.hs view
@@ -1,12 +1,9 @@-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE PolyKinds #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE TypeOperators #-} -{- | Likelihood-Weighting inference+{- | Likelihood-Weighting inference. -}  module Inference.LW (@@ -29,28 +26,36 @@  -- | Top-level wrapper for Likelihood-Weighting (LW) inference lw :: (FromSTrace env, es ~ '[ObsReader env, Dist, State STrace, Observe, Sample])-    => -    -- | Number of LW iterations-       Int                          -    -- | Model awaiting an input-    -> (b -> Model env es a)        -    -- | Model input and model environment (containing observed values to condition on)-    -> (b, Env env)                 -    -- | Trace of weighted output environments containing values sampled for each LW iteration-    -> Sampler [(Env env, Double)]  +    -- | number of LW iterations+    => Int+    -- | model awaiting an input+    -> (b -> Model env es a)+    -- | (model input, input model environment)+    -> (b, Env env)+    -- | [(output model environment, likelihood-weighting)]+    -> Sampler [(Env env, Double)] lw n model xs_envs = do   let runN (x, env) = replicateM n (runLW env (model x))   lwTrace <- runN xs_envs   return $ map (\((_, strace), p) -> (fromSTrace strace, p)) lwTrace  -- | Handler for one iteration of LW-runLW :: es ~ '[ObsReader env, Dist,State STrace,  Observe, Sample]-  => Env env -> Model env es a+runLW :: es ~ '[ObsReader env, Dist, State STrace, Observe, Sample]+  -- | model environment+  => Env env+  -- | model+  -> Model env es a+  -- | ((model output, sample trace), likelihood-weighting)   -> Sampler ((a, STrace), Double) runLW env = handleSamp . handleObs 0 . handleState Map.empty . traceSamples . handleCore env --- | Handle each Observe operation by computing and accumulating log probabilities-handleObs :: Member Sample es => Double -> Prog (Observe : es) a -> Prog es (a, Double)+-- | Handle each @Observe@ operation by computing and accumulating a log probability+handleObs :: Member Sample es+  -- | accumulated log-probability+  => Double+  -> Prog (Observe : es) a+  -- | (model output, final log-probability)+  -> Prog es (a, Double) handleObs logp (Val x) = return (x, exp logp) handleObs logp (Op u k) = case discharge u of     Right (Observe d y α) -> do
src/Inference/MH.hs view
@@ -1,15 +1,13 @@ {-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE TypeOperators #-} {-# LANGUAGE GADTs #-}-{-# OPTIONS_GHC -Wno-incomplete-patterns #-} {-# LANGUAGE PatternSynonyms #-} {-# LANGUAGE ViewPatterns #-} -{- | Metropolis-Hastings inference +{- | Metropolis-Hastings inference. -} -module Inference.MH ( +module Inference.MH (     mh   , mhStep   , runMH@@ -31,6 +29,11 @@ import Model ( Model, handleCore ) import OpenSum (OpenSum(..)) import PrimDist+    ( ErasedPrimDist(ErasedPrimDist),+      PrimVal,+      PrimDist(UniformDist, DiscrUniformDist),+      pattern PrimDistPrf,+      sample ) import Prog ( Member(prj), EffectSum, Prog(..), discharge ) import qualified Data.Map as Map import qualified Data.Set as Set@@ -39,19 +42,19 @@ import Trace ( LPTrace, FromSTrace(..), STrace, updateLPTrace ) import Unsafe.Coerce ( unsafeCoerce ) --- | Top-level wrapper for Metropolis-Hastings (MH) inference +-- | Top-level wrapper for Metropolis-Hastings (MH) inference mh :: (FromSTrace env, es ~ '[ObsReader env, Dist, State STrace, State LPTrace, Observe, Sample])-  => -  -- | Number of MH iterations-     Int                    -  -- | A model awaiting an input-  -> (b -> Model env es a)  -  -- | A model input and model environment (containing observed values to condition on)-  -> (b, Env env)           -  -- | An optional list of observable variable names (strings) to specify sample sites of interest (e.g. for interest in sampling #mu, provide "mu"). This causes other variables to not be resampled unless necessary.-  -> [Tag]                  -  -- | Trace of output environments, containing values sampled for each MH iteration-  -> Sampler [Env env]     +  -- | number of MH iterations+  => Int+  -- | model awaiting an input+  -> (b -> Model env es a)+  -- | (model input, input model environment)+  -> (b, Env env)+  -- | optional list of observable variable names (strings) to specify sample sites of interest+  {- For example, provide "mu" to specify interest in sampling #mu. This causes other variables to not be resampled unless necessary. -}+  -> [Tag]+  -- | [output model environment]+  -> Sampler [Env env] mh n model  (x_0, env_0) tags = do   -- Perform initial run of MH with no proposal sample site   y0 <- runMH env_0 Map.empty ("", 0) (model x_0)@@ -62,48 +65,45 @@  -- | Perform one step of MH mhStep :: (es ~ '[ObsReader env, Dist, State STrace, State LPTrace, Observe, Sample])-  => -  -- | Model environment-     Env env                  -  -- | Model-  -> Model env es a            -  -- | Tags indicating sample sites of interest-  -> [Tag]                     -  -- | Trace of previous MH outputs-  -> [((a, STrace), LPTrace)]  -  -- | Updated trace of MH outputs+  -- | model environment+  => Env env+  -- | model+  -> Model env es a+  -- | tags indicating sample sites of interest+  -> [Tag]+  -- | trace of previous MH outputs+  -> [((a, STrace), LPTrace)]+  -- | updated trace of MH outputs   -> Sampler [((a, STrace), LPTrace)] mhStep env model tags trace = do-  let -- Get previous mh output-      ((x, samples), logps) = head trace-+  -- Get previous mh output+  let ((x, samples), logps) = head trace+  -- Get possible addresses to propose new samples for       sampleSites = if null tags then samples                     else  Map.filterWithKey (\(tag, i) _ -> tag `elem` tags) samples-+  -- Draw a proposal sample address   α_samp_ind <- sample $ DiscrUniformDist 0 (Map.size sampleSites - 1)   let (α_samp, _) = Map.elemAt α_samp_ind sampleSites-+  -- Run MH with proposal sample address   ((x', samples'), logps') <- runMH env samples α_samp model-+  -- Compute acceptance ratio   acceptance_ratio <- liftS $ accept α_samp samples samples' logps logps'   u <- sample (UniformDist 0 1)-   if u < acceptance_ratio     then do return (((x', samples'), logps'):trace)     else do return trace  -- | Handler for one iteration of MH runMH :: (es ~ '[ObsReader env, Dist, State STrace, State LPTrace, Observe, Sample])-  => -  -- | Model environment-     Env env       -  -- | Sample trace of previous MH iteration-  -> STrace         -  -- | Sample address of interest-  -> Addr           -  -- | Model+  -- | model environment+  => Env env+  -- | sample trace of previous MH iteration+  -> STrace+  -- | sample address of interest+  -> Addr+  -- | model   -> Model env es a-  -- | (Outpiut, sample trace, log)+  -- | (model output, sample trace, log-probability trace)   -> Sampler ((a, STrace), LPTrace) runMH env strace α_samp =      handleSamp strace α_samp  . handleObs@@ -115,9 +115,11 @@  pattern Obs :: Member Observe es => PrimDist x -> x -> Addr -> EffectSum es x pattern Obs d y α <- (prj -> Just (Observe d y α))- --- | Handler for tracing log-probabilities for each Sample and Observe operation-traceLPs :: (Member (State LPTrace) es, Member Sample es, Member Observe es) => Prog es a -> Prog es a++-- | Handler for tracing log-probabilities for each @Sample@ and @Observe@ operation+traceLPs :: (Member (State LPTrace) es, Member Sample es, Member Observe es)+  => Prog es a+  -> Prog es a traceLPs (Val x) = return x traceLPs (Op op k) = case op of   Samp (PrimDistPrf d) α ->@@ -128,14 +130,13 @@                   traceLPs (k y))   _         -> Op op (traceLPs . k) --- | Handler for Sample that selectively reuses old samples or draws new ones-handleSamp :: -  -- | Sample trace-     STrace -  -- | Address of the proposal sample site for the current MH iteration-  -> Addr -  -- |  Idx                                                                                                                         -  -> Prog '[Sample] a +-- | Handler for @Sample@ that selectively reuses old samples or draws new ones+handleSamp ::+  -- | sample trace+     STrace+  -- | address of the proposal sample site for the current MH iteration+  -> Addr+  -> Prog '[Sample] a   -> Sampler a handleSamp strace α_samp (Op op k) = case discharge op of   Right (Sample (PrimDistPrf d) α) ->@@ -146,16 +147,16 @@  -- | For a given address, look up a sampled value from a sample trace, returning --   it only if the primitive distribution it was sampled from matches the current one.-lookupSample :: OpenSum.Member a PrimVal -  => -  -- | Sample trace-     STrace -  -- | Distribution to sample from-  -> PrimDist a -  -- | Address of current sample site-  -> Addr -  -- | Address of proposal sample site-  -> Addr +lookupSample :: OpenSum.Member a PrimVal+  =>+  -- | sample trace+     STrace+  -- | distribution to sample from+  -> PrimDist a+  -- | address of current sample site+  -> Addr+  -- | address of proposal sample site+  -> Addr   -> Sampler a lookupSample samples d α α_samp   | α == α_samp = sample d@@ -168,17 +169,17 @@         Nothing -> sample d  -- | Compute acceptance probability-accept :: -  -- | Address of new sampled value-     Addr -  -- | Previous MH sample trace+accept ::+  -- | address of new sampled value+     Addr+  -- | previous MH sample trace   -> STrace-  -- | New MH sample trace-  -> STrace -  -- | Previous MH log-probability trace-  -> LPTrace -  -- | Current MH log-probability trace-  -> LPTrace +  -- | new MH sample trace+  -> STrace+  -- | previous MH log-probability trace+  -> LPTrace+  -- | current MH log-probability trace+  -> LPTrace   -> IO Double accept x0 _Ⲭ _Ⲭ' logℙ logℙ' = do   let _X'sampled = Set.singleton x0 `Set.union` (Map.keysSet _Ⲭ' \\ Map.keysSet _Ⲭ)
src/Inference/SIM.hs view
@@ -1,12 +1,12 @@-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-} {-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-}+{-# LANGUAGE PatternSynonyms #-}+{-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeOperators #-}-{-# LANGUAGE FlexibleContexts #-} {-# OPTIONS_GHC -Wno-incomplete-patterns #-} -{- | Simulation +{- | Simulation. -}  module Inference.SIM (@@ -23,7 +23,7 @@ import Env ( Env ) import Model ( Model, handleCore ) import OpenSum (OpenSum)-import PrimDist+import PrimDist ( pattern PrimDistPrf, sample ) import Prog ( Member(prj), Prog(..), discharge ) import qualified Data.Map as Map import qualified OpenSum@@ -33,29 +33,33 @@  -- | Top-level wrapper for simulating from a model simulate :: (FromSTrace env, es ~ '[ObsReader env, Dist,State STrace, Observe, Sample])-  => -  -- | A model awaiting an input-  (b -> Model env es a)  -  -- | A model environment-  -> Env env               -  -- | Model input -  -> b                    -  -- | Model output and output environment  -  -> Sampler (a, Env env)   +  -- | model awaiting an input+  => (b -> Model env es a)+  -- | model environment+  -> Env env+  -- | model input+  -> b+  -- | (model output, output environment)+  -> Sampler (a, Env env) simulate model env x  = do   outputs_strace <- runSimulate env (model x)   return (fmap fromSTrace outputs_strace)  -- | Handler for simulating once from a probabilistic program runSimulate :: (es ~ '[ObsReader env, Dist, State STrace, Observe, Sample])- => Env env -> Model env es a -> Sampler (a, STrace)+ -- | model environment+ => Env env+ -- | model+ -> Model env es a+ -- | (model output, sample trace)+ -> Sampler (a, STrace) runSimulate env   = handleSamp . handleObs . handleState Map.empty . traceSamples . handleCore env --- | Trace sampled values for each Sample operation+-- | Trace sampled values for each @Sample@ operation traceSamples :: (Member (State STrace) es, Member Sample es) => Prog es a -> Prog es a traceSamples (Val x) = return x-traceSamples (Op op k) = case prj op of +traceSamples (Op op k) = case prj op of   Just (Sample (PrimDistPrf d) α) ->        Op op (\x -> do modify (updateSTrace α d x);                        traceSamples (k x))
src/Model.hs view
@@ -1,17 +1,16 @@-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE DeriveFunctor #-}-{-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE DataKinds #-}+{-# LANGUAGE DeriveFunctor #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE MonoLocalBinds #-} {-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeOperators #-}-{-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE MonoLocalBinds #-}  {- | An algebraic effect embedding of probabilistic models. -} -module Model ( +module Model (     Model(..)   , handleCore     -- * Distribution smart constructors@@ -44,7 +43,7 @@   , poisson'   , uniform   , uniform'-  ) +  )   where  import Control.Monad ( ap )@@ -54,20 +53,20 @@ import Effects.ObsReader ( ask, handleRead, ObsReader ) import Effects.State ( State, modify, handleState ) import Env ( varToStr, Env, ObsVar, Observable )-import OpenSum (OpenSum)-import PrimDist+import PrimDist ( PrimVal, PrimDist(..) ) import Prog ( call, Member, Prog ) import qualified OpenSum  {- | Models are parameterised by: -1) a model environment @env@ containing random variables that can be provided observed values for, +    1) a model environment @env@ containing random variables that can be provided observed values -2) an effect signature @es@ of the possible effects a model can invoke, and +    2) an effect signature @es@ of the possible effects a model can invoke -3) an output type @a@ of values that the model generates. +    3) an output type @a@ of values that the model generates. -A model initially consists of (at least) two effects: @Dist@ for calling primitive distributions and @ObsReader env@ for reading from @env@. +    A model initially consists of (at least) two effects: @Dist@ for calling primitive distributions+    and @ObsReader env@ for reading from @env@. -} newtype Model env es a =   Model { runModel :: (Member Dist es, Member (ObsReader env) es) => Prog es a }@@ -83,51 +82,51 @@     f' <- f     runModel $ x f' -{- | The initial handler for models, specialising a model under a certain +{- | The initial handler for models, specialising a model under a certain environment to produce a probabilistic program consisting of @Sample@ and @Observe@ operations. -} handleCore :: (Member Observe es, Member Sample es) => Env env -> Model env (ObsReader env : Dist : es) a -> Prog es a handleCore env m = (handleDist . handleRead env) (runModel m)  {- $Smart-Constructors -Smart constructors for calling primitive distribution operations inside models, -where each distribution comes with a primed and an unprimed variant.--An unprimed distribution takes the standard distribution parameters as well as -an observable variable. This lets one later provide observed values for that -variable to be conditioned against:+    Smart constructors for calling primitive distribution operations inside models,+    where each distribution comes with a primed and an unprimed variant. -@-exampleModel :: Observable env "b" Bool => Model env es Bool-exampleModel = bernoulli 0.5 #b-@+    An unprimed distribution takes the standard distribution parameters as well as+    an observable variable. This lets one later provide observed values for that+    variable to be conditioned against: -A primed distribution takes no observable variable and so cannot be conditioned against; this will always representing sampling from that distribution:+    @+    exampleModel :: Observable env "b" Bool => Model env es Bool+    exampleModel = bernoulli 0.5 #b+    @ -@-exampleModel' :: Model env es Bool-exampleModel' = bernoulli' 0.5-@+    A primed distribution takes no observable variable and so cannot be conditioned against;+    this will always representing sampling from that distribution: +    @+    exampleModel' :: Model env es Bool+    exampleModel' = bernoulli' 0.5+    @ -} -deterministic :: forall env es a x. (Eq a, Show a, OpenSum.Member a PrimVal, Observable env x a) => a    -  -> ObsVar x +deterministic :: forall env es a x. (Eq a, Show a, OpenSum.Member a PrimVal, Observable env x a) => a+  -> ObsVar x   -> Model env es a deterministic x field = Model $ do   let tag = Just $ varToStr field   maybe_y <- ask @env field   call (Dist (DeterministicDist x) maybe_y tag) -deterministic' :: (Eq a, Show a, OpenSum.Member a PrimVal) => -  -- | Value to be deterministically generated-     a +deterministic' :: (Eq a, Show a, OpenSum.Member a PrimVal) =>+  -- | value to be deterministically generated+     a   -> Model env es a deterministic' x = Model $ do   call (Dist (DeterministicDist x) Nothing Nothing) -dirichlet :: forall env es x. Observable env x [Double] => -     [Double] +dirichlet :: forall env es x. Observable env x [Double] =>+     [Double]   -> ObsVar x   -> Model env es [Double] dirichlet xs field = Model $ do@@ -135,15 +134,15 @@   maybe_y <- ask @env field   call (Dist (DirichletDist xs) maybe_y tag) -dirichlet' :: -  -- | Concentration parameters-     [Double] +dirichlet' ::+  -- | concentration parameters+     [Double]   -> Model env es [Double] dirichlet' xs = Model $ do   call (Dist (DirichletDist xs) Nothing Nothing) -discrete :: forall env es x. Observable env x Int => -     [Double] +discrete :: forall env es x. Observable env x Int =>+     [Double]   -> ObsVar x   -> Model env es Int discrete ps field = Model $ do@@ -151,16 +150,16 @@   maybe_y <- ask @env field   call (Dist (DiscreteDist ps) maybe_y tag) -discrete' :: -  -- | List of @n@ probabilities-     [Double] -  -- | Integer index from @0@ to @n - 1@+discrete' ::+  -- | list of @n@ probabilities+     [Double]+  -- | integer index from @0@ to @n - 1@   -> Model env es Int discrete' ps = Model $ do   call (Dist (DiscreteDist ps) Nothing Nothing) -categorical :: forall env es a x. (Eq a, Show a, OpenSum.Member a PrimVal, Observable env x a) => -     [(a, Double)] +categorical :: forall env es a x. (Eq a, Show a, OpenSum.Member a PrimVal, Observable env x a) =>+     [(a, Double)]   -> ObsVar x   -> Model env es a categorical xs field = Model $ do@@ -168,16 +167,16 @@   maybe_y <- ask @env field   call (Dist (CategoricalDist xs) maybe_y tag) -categorical' :: (Eq a, Show a, OpenSum.Member a PrimVal) => -  -- | Primitive values and their probabilities-     [(a, Double)] +categorical' :: (Eq a, Show a, OpenSum.Member a PrimVal) =>+  -- | primitive values and their probabilities+     [(a, Double)]   -> Model env es a categorical' xs = Model $ do   call (Dist (CategoricalDist xs) Nothing Nothing) -normal :: forall env es x. Observable env x Double => +normal :: forall env es x. Observable env x Double =>      Double-  -> Double +  -> Double   -> ObsVar x   -> Model env es Double normal mu sigma field = Model $ do@@ -185,17 +184,17 @@   maybe_y <- ask @env field   call (Dist (NormalDist mu sigma) maybe_y tag) -normal' :: -  -- | Mean-     Double -  -- | Standard deviation -  -> Double +normal' ::+  -- | mean+     Double+  -- | standard deviation+  -> Double   -> Model env es Double normal' mu sigma = Model $ do   call (Dist (NormalDist mu sigma) Nothing Nothing) -halfNormal :: forall env es x. Observable env x Double => -     Double +halfNormal :: forall env es x. Observable env x Double =>+     Double   -> ObsVar x   -> Model env es Double halfNormal sigma field = Model $ do@@ -203,16 +202,16 @@   maybe_y <- ask @env field   call (Dist (HalfNormalDist sigma) maybe_y tag) -halfNormal' :: -  -- | Standard deviation-     Double +halfNormal' ::+  -- | standard deviation+     Double   -> Model env es Double halfNormal' sigma = Model $ do   call (Dist (HalfNormalDist sigma) Nothing Nothing) -cauchy :: forall env es x. Observable env x Double => -     Double -  -> Double +cauchy :: forall env es x. Observable env x Double =>+     Double+  -> Double   -> ObsVar x   -> Model env es Double cauchy mu sigma field = Model $ do@@ -220,17 +219,17 @@   maybe_y <- ask @env field   call (Dist (CauchyDist mu sigma) maybe_y tag) -cauchy' :: -  -- | Location-     Double -  -- | Scale-  -> Double +cauchy' ::+  -- | location+     Double+  -- | scale+  -> Double   -> Model env es Double cauchy' mu sigma = Model $ do   call (Dist (CauchyDist mu sigma) Nothing Nothing) -halfCauchy :: forall env es x. Observable env x Double => -     Double +halfCauchy :: forall env es x. Observable env x Double =>+     Double   -> ObsVar x   -> Model env es Double halfCauchy sigma field = Model $ do@@ -238,15 +237,15 @@   maybe_y <- ask @env field   call (Dist (HalfCauchyDist sigma) maybe_y tag) -halfCauchy' :: -  -- | Scale-     Double +halfCauchy' ::+  -- | scale+     Double   -> Model env es Double halfCauchy' sigma = Model $ do   call (Dist (HalfCauchyDist sigma) Nothing Nothing) -bernoulli :: forall env es x. Observable env x Bool => -     Double +bernoulli :: forall env es x. Observable env x Bool =>+     Double   -> ObsVar x   -> Model env es Bool bernoulli p field = Model $ do@@ -254,16 +253,16 @@   maybe_y <- ask @env field   call (Dist (BernoulliDist p) maybe_y tag) -bernoulli' :: -  -- | Probability of @True@-     Double +bernoulli' ::+  -- | probability of @True@+     Double   -> Model env es Bool bernoulli' p = Model $ do   call (Dist (BernoulliDist p) Nothing Nothing) -beta :: forall env es x. Observable env x Double => -     Double -  -> Double +beta :: forall env es x. Observable env x Double =>+     Double+  -> Double   -> ObsVar x   -> Model env es Double beta α β field = Model $ do@@ -271,18 +270,18 @@   maybe_y <- ask @env field   call (Dist (BetaDist α β) maybe_y tag) -beta' :: -  -- | Shape 1 (α)+beta' ::+  -- | shape 1 (α)      Double-  -- | Shape 2 (β)-  -> Double +  -- | shape 2 (β)+  -> Double   -> Model env es Double beta' α β = Model $ do   call (Dist (BetaDist α β) Nothing Nothing) -binomial :: forall env es x. Observable env x Int => -     Int -  -> Double +binomial :: forall env es x. Observable env x Int =>+     Int+  -> Double   -> ObsVar x   -> Model env es Int binomial n p field = Model $ do@@ -290,19 +289,19 @@   maybe_y <- ask @env field   call (Dist (BinomialDist n p) maybe_y tag) -binomial' :: -  -- | Number of trials   -     Int -  -- | Probability of successful trial-  -> Double -  -- | Number of successful trials+binomial' ::+  -- | number of trials+     Int+  -- | probability of successful trial+  -> Double+  -- | number of successful trials   -> Model env es Int binomial' n p = Model $ do   call (Dist (BinomialDist n p) Nothing Nothing) -gamma :: forall env es x. Observable env x Double => -     Double -  -> Double +gamma :: forall env es x. Observable env x Double =>+     Double+  -> Double   -> ObsVar x   -> Model env es Double gamma k θ field = Model $ do@@ -310,18 +309,18 @@   maybe_y <- ask @env field   call (Dist (GammaDist k θ) maybe_y tag) -gamma' :: -  -- | Shape (k)-     Double -  -- | Scale (θ)-  -> Double +gamma' ::+  -- | shape (k)+     Double+  -- | scale (θ)+  -> Double   -> Model env es Double gamma' k θ = Model $ do   call (Dist (GammaDist k θ) Nothing Nothing) -uniform :: forall env es x. Observable env x Double => -     Double -  -> Double +uniform :: forall env es x. Observable env x Double =>+     Double+  -> Double   -> ObsVar x   -> Model env es Double uniform min max field = Model $ do@@ -330,16 +329,16 @@   call (Dist (UniformDist min max) maybe_y tag)  uniform' ::-  -- | Lower-bound-     Double -  -- | Upper-bound-  -> Double +  -- | lower-bound+     Double+  -- | upper-bound+  -> Double   -> Model env es Double uniform' min max = Model $ do   call (Dist (UniformDist min max) Nothing Nothing) -poisson :: forall env es x. Observable env x Int => -     Double +poisson :: forall env es x. Observable env x Int =>+     Double   -> ObsVar x   -> Model env es Int poisson λ field = Model $ do@@ -347,10 +346,10 @@   maybe_y <- ask @env field   call (Dist (PoissonDist λ) maybe_y tag) -poisson' :: -  -- | Rate (λ)-     Double -  -- | Number of events+poisson' ::+  -- | rate (λ)+     Double+  -- | number of events   -> Model env es Int poisson' λ = Model $ do   call (Dist (PoissonDist λ) Nothing Nothing)
src/OpenSum.hs view
@@ -1,18 +1,16 @@-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE EmptyCase #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleInstances #-}+{-# LANGUAGE GADTs #-} {-# LANGUAGE MultiParamTypeClasses #-}+{-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilyDependencies #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE DataKinds #-}-{-# LANGUAGE PolyKinds #-} {-# LANGUAGE TypeOperators #-}-{-# LANGUAGE GADTs #-}-{-# LANGUAGE FlexibleInstances #-}-{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE EmptyCase #-} {-# LANGUAGE UndecidableInstances #-}  {- | An open sum implementation for value types. -}+ module OpenSum (     OpenSum(..)   , Member(..)) where
src/PrimDist.hs view
@@ -1,23 +1,11 @@-{-# LANGUAGE PatternSynonyms #-}-{-# LANGUAGE RankNTypes #-}-{-# LANGUAGE GADTs, TypeOperators #-}-{-# LANGUAGE TypeApplications #-} {-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-}-{-# LANGUAGE KindSignatures #-}-{-# LANGUAGE PolyKinds #-}-{-# LANGUAGE TypeFamilies #-}-{-# LANGUAGE ViewPatterns #-}-{-# LANGUAGE FlexibleInstances #-}-{-# LANGUAGE UndecidableInstances #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE PatternSynonyms #-} {-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE LambdaCase #-}-{-# LANGUAGE EmptyCase #-}-{-# LANGUAGE ConstraintKinds #-}-{-# LANGUAGE ImplicitParams #-}-{-# OPTIONS_GHC -Wno-incomplete-patterns #-}+{-# LANGUAGE ViewPatterns #-} -{- | A GADT encoding of (a selection of) primitive distributions +{- | A GADT encoding of (a selection of) primitive distributions     along with their corresponding sampling and density functions. -} @@ -35,12 +23,9 @@   , logProb) where  import Data.Kind ( Constraint )-import Data.Map (Map)-import Numeric.Log ( Log(Exp) )-import OpenSum (OpenSum)+import Numeric.Log ( Log(..) ) import qualified Data.Map as Map import qualified Data.Vector as V-import qualified Data.Vector as Vec import qualified Data.Vector.Unboxed as UV import qualified OpenSum import qualified System.Random.MWC.Distributions as MWC@@ -59,60 +44,60 @@  -- | Primitive distribution data PrimDist a where-  BernoulliDist     -    :: Double           -- ^ Probability of @True@-    -> PrimDist Bool  -  BetaDist          -    :: Double           -- ^ Shape α-    -> Double           -- ^ Shape β+  BernoulliDist+    :: Double           -- ^ probability of @True@+    -> PrimDist Bool+  BetaDist+    :: Double           -- ^ shape α+    -> Double           -- ^ shape β     -> PrimDist Double-  BinomialDist      -    :: Int              -- ^ Number of trials-    -> Double           -- ^ Probability of successful trial-    -> PrimDist Int     -  CategoricalDist   -    :: (Eq a, Show a, OpenSum.Member a PrimVal) -    => [(a, Double)]    -- ^ Values and associated probabilities-    -> PrimDist a       -  CauchyDist        -    :: Double           -- ^ Location-    -> Double           -- ^ Scale+  BinomialDist+    :: Int              -- ^ number of trials+    -> Double           -- ^ probability of successful trial+    -> PrimDist Int+  CategoricalDist+    :: (Eq a, Show a, OpenSum.Member a PrimVal)+    => [(a, Double)]    -- ^ values and associated probabilities+    -> PrimDist a+  CauchyDist+    :: Double           -- ^ location+    -> Double           -- ^ scale     -> PrimDist Double-  HalfCauchyDist      -    :: Double           -- ^ Scale+  HalfCauchyDist+    :: Double           -- ^ scale     -> PrimDist Double-  DeterministicDist -    :: (Eq a, Show a, OpenSum.Member a PrimVal) -    => a                -- ^ Value of probability @1@+  DeterministicDist+    :: (Eq a, Show a, OpenSum.Member a PrimVal)+    => a                -- ^ value of probability @1@     -> PrimDist a-  DirichletDist     -    :: [Double]         -- ^ Concentrations+  DirichletDist+    :: [Double]         -- ^ concentrations     -> PrimDist [Double]-  DiscreteDist      -    :: [Double]         -- ^ List of @n@ probabilities-    -> PrimDist Int     -- ^ An index from @0@ to @n - 1@-  DiscrUniformDist      -    :: Int              -- ^ Lower-bound @a@-    -> Int              -- ^ Upper-bound @b@-    -> PrimDist Int     -  GammaDist         -    :: Double           -- ^ Shape k-    -> Double           -- ^ Scale θ+  DiscreteDist+    :: [Double]         -- ^ list of @n@ probabilities+    -> PrimDist Int     -- ^ an index from @0@ to @n - 1@+  DiscrUniformDist+    :: Int              -- ^ lower-bound @a@+    -> Int              -- ^ upper-bound @b@+    -> PrimDist Int+  GammaDist+    :: Double           -- ^ shape k+    -> Double           -- ^ scale θ     -> PrimDist Double-  NormalDist      -    :: Double           -- ^ Mean-    -> Double           -- ^ Standard deviation+  NormalDist+    :: Double           -- ^ mean+    -> Double           -- ^ standard deviation     -> PrimDist Double-  HalfNormalDist    -    :: Double           -- ^ Standard deviation+  HalfNormalDist+    :: Double           -- ^ standard deviation     -> PrimDist Double-  PoissonDist       -    :: Double           -- ^ Rate λ+  PoissonDist+    :: Double           -- ^ rate λ     -> PrimDist Int-  UniformDist       -    :: Double           -- ^ Lower-bound @a@-    -> Double           -- ^ Upper-bound @b@-    -> PrimDist Double  +  UniformDist+    :: Double           -- ^ lower-bound @a@+    -> Double           -- ^ upper-bound @b@+    -> PrimDist Double  instance Eq (PrimDist a) where   (==) (NormalDist m s) (NormalDist m' s') = m == m' && s == s'@@ -163,7 +148,7 @@    "DirichletDist(" ++ show xs ++ ", " ++ ")"   show (DeterministicDist x) =    "DeterministicDist(" ++ show x ++ ", " ++ ")"-  + -- | An ad-hoc specification of primitive value types, for constraining the outputs of distributions type PrimVal = '[Int, Double, [Double], Bool, String] @@ -171,12 +156,12 @@ data IsPrimVal x where   IsPrimVal :: (Show x, OpenSum.Member x PrimVal) => IsPrimVal x --- | For pattern-matching on an arbitrary @PrimDist@ with proof that it generates a primitive value +-- | For pattern-matching on an arbitrary @PrimDist@ with proof that it generates a primitive value pattern PrimDistPrf :: () => (Show x, OpenSum.Member x PrimVal) => PrimDist x -> PrimDist x pattern PrimDistPrf d <- d@(primDistPrf -> IsPrimVal)  -- | Proof that all primitive distributions generate a primitive value-primDistPrf :: PrimDist x -> IsPrimVal x +primDistPrf :: PrimDist x -> IsPrimVal x primDistPrf d = case d of   HalfCauchyDist {} -> IsPrimVal   CauchyDist {} -> IsPrimVal@@ -202,8 +187,8 @@   show (ErasedPrimDist d) = show d  -- | Draw a value from a primitive distribution in the @Sampler@ monad-sample :: -     PrimDist a +sample ::+     PrimDist a   -> Sampler a sample (HalfCauchyDist σ )  =   createSampler (sampleCauchy 0 σ) >>= pure . abs@@ -236,12 +221,12 @@ sample (DeterministicDist x) = pure x  -- | Compute the density of a primitive distribution generating an observed value-prob :: -  -- | Distribution-     PrimDist a -  -- | Observed value-  -> a -  -- | Density+prob ::+  -- | distribution+     PrimDist a+  -- | observed value+  -> a+  -- | density   -> Double prob (DirichletDist xs) ys =   let xs' = map (/(Prelude.sum xs)) xs@@ -276,19 +261,16 @@   = case lookup y ps of       Nothing -> error $ "Couldn't find " ++ show y ++ " in categorical dist"       Just p  -> p-prob (DiscreteDist ps) y-  = ps !! y-prob (PoissonDist λ) y-  = probability (poisson λ) y-prob (DeterministicDist x) y-  = 1+prob (DiscreteDist ps) y     = ps !! y+prob (PoissonDist λ) y       = probability (poisson λ) y+prob (DeterministicDist x) y = 1  -- | Compute the log density of a primitive distribution generating an observed value-logProb :: -  -- | Distribution-     PrimDist a -  -- | Observed value-  -> a -  -- | Log density+logProb ::+  -- | distribution+     PrimDist a+  -- | observed value+  -> a+  -- | log density   -> Double logProb d = log . prob d
src/Prog.hs view
@@ -1,15 +1,14 @@ {-# LANGUAGE DataKinds #-}+{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE MultiParamTypeClasses #-}-{-# LANGUAGE PolyKinds #-} {-# LANGUAGE PatternSynonyms, ViewPatterns #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilyDependencies #-}-{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE TypeOperators #-} {-# LANGUAGE UndecidableInstances #-} -{- | An encoding for algebraic effects, based on the @freer@ monad. +{- | An encoding for algebraic effects, based on the @freer@ monad. -}  module Prog (@@ -28,14 +27,16 @@ import GHC.TypeLits ( TypeError, ErrorMessage(Text, (:<>:), (:$$:), ShowType) ) import Unsafe.Coerce ( unsafeCoerce ) --- | A program that returns a value of type @a@ and can call operations that belong to some effect @e@ in signature @es@; this represents a syntax tree whose nodes are operations and leaves are pure values.+{- | A program that returns a value of type @a@ and can call operations that belong to some effect+     @e@ in signature @es@; this represents a syntax tree whose nodes are operations and leaves are pure values.+-} data Prog es a where-  Val -    :: a                -- ^ pure value +  Val+    :: a                -- ^ pure value     -> Prog es a-  Op -    :: EffectSum es x   -- ^ an operation belonging to some effect in @es@-    -> (x -> Prog es a) -- ^ a continuation from the result of the operation+  Op+    :: EffectSum es x   -- ^ operation belonging to some effect in @es@+    -> (x -> Prog es a) -- ^ continuation from the result of the operation     -> Prog es a  instance Functor (Prog es) where@@ -59,9 +60,9 @@ -- | Membership of an effect @e@ in @es@ class (FindElem e es) => Member (e :: * -> *) (es :: [* -> *]) where   -- | Inject an operation of type @e x@ into an effect sum-  inj ::  e x -> EffectSum es x+  inj :: e x -> EffectSum es x   -- | Attempt to project an operation of type @e x@ out from an effect sum-  prj ::  EffectSum es x -> Maybe (e x)+  prj :: EffectSum es x -> Maybe (e x)  instance {-# INCOHERENT #-} (e ~ e') => Member e '[e'] where    inj x  = EffectSum 0 x
src/Sampler.hs view
@@ -1,6 +1,6 @@ {-# LANGUAGE GeneralizedNewtypeDeriving #-} -{- | An IO-based sampling monad+{- | An IO-based sampling monad. -}  module Sampler (@@ -28,14 +28,10 @@   ) where  import Control.Monad ( replicateM )-import Control.Monad.ST (ST, runST, stToIO) import Control.Monad.Trans (MonadIO, MonadTrans, lift) import Control.Monad.Trans.Reader (ReaderT, ask, mapReaderT, runReaderT) import Data.Map (Map)-import Data.Set (Set) import GHC.Word ( Word32 )-import qualified Data.Map as Map-import qualified Data.Set as Set import qualified Data.Vector as V import qualified System.Random.MWC as MWC import qualified System.Random.MWC.Distributions as MWC.Dist@@ -69,77 +65,78 @@ createSampler f = Sampler $ ask >>= lift . f  {- $Sampling-functions-Given their distribution parameters, these functions await a generator and then sample a value from the distribution in the @IO@ monad+  Given their distribution parameters, these functions await a generator and+  then sample a value from the distribution in the @IO@ monad. -} -sampleRandom -  :: MWC.GenIO +sampleRandom+  :: MWC.GenIO   -> IO Double sampleRandom = \gen -> MWC.uniform gen -sampleCauchy -  :: Double -- ^ Location-  -> Double -- ^ Scale+sampleCauchy+  :: Double -- ^ location+  -> Double -- ^ scale   -> (MWC.GenIO -> IO Double) sampleCauchy μ σ = \gen -> genContVar (cauchyDistribution μ σ) gen -sampleNormal -  :: Double -- ^ Mean-  -> Double -- ^ Standard deviation+sampleNormal+  :: Double -- ^ mean+  -> Double -- ^ standard deviation   -> (MWC.GenIO -> IO Double) sampleNormal μ σ = \gen -> MWC.Dist.normal μ σ gen -sampleUniform -  :: Double -- ^ Lower-bound-  -> Double -- ^ Upper-bound+sampleUniform+  :: Double -- ^ lower-bound+  -> Double -- ^ upper-bound   -> (MWC.GenIO -> IO Double) sampleUniform min max = \gen -> MWC.uniformR (min, max) gen -sampleDiscreteUniform -  :: Int -- ^ Lower-bound-  -> Int -- ^ Upper-bound+sampleDiscreteUniform+  :: Int -- ^ lower-bound+  -> Int -- ^ upper-bound   -> (MWC.GenIO -> IO Int) sampleDiscreteUniform min max = \gen -> MWC.uniformR (min, max) gen -sampleGamma -  :: Double -- ^ Shape k-  -> Double -- ^ Scale θ+sampleGamma+  :: Double -- ^ shape k+  -> Double -- ^ scale θ   -> (MWC.GenIO -> IO Double) sampleGamma k θ = \gen -> MWC.Dist.gamma k θ gen -sampleBeta -  :: Double -- ^ Shape α-  -> Double -- ^ Shape β+sampleBeta+  :: Double -- ^ shape α+  -> Double -- ^ shape β   -> (MWC.GenIO -> IO Double) sampleBeta α β = \gen -> MWC.Dist.beta α β gen -sampleBernoulli -  :: Double -- ^ Probability of @True@+sampleBernoulli+  :: Double -- ^ probability of @True@   -> (MWC.GenIO -> IO Bool) sampleBernoulli p = \gen -> MWC.Dist.bernoulli p gen -sampleBinomial -  :: Int    -- ^ Number of trials-  -> Double -- ^ Probability of successful trial+sampleBinomial+  :: Int    -- ^ number of trials+  -> Double -- ^ probability of successful trial   -> (MWC.GenIO -> IO [Bool]) sampleBinomial n p = \gen -> replicateM n (MWC.Dist.bernoulli p gen) -sampleCategorical -  :: V.Vector Double -- ^ Probabilities+sampleCategorical+  :: V.Vector Double -- ^ probabilities   -> (MWC.GenIO -> IO Int) sampleCategorical ps =  \gen -> MWC.Dist.categorical (ps) gen -sampleDiscrete -  :: [Double] -- ^ Probabilities+sampleDiscrete+  :: [Double] -- ^ probabilities   -> (MWC.GenIO -> IO Int) sampleDiscrete ps = \gen -> MWC.Dist.categorical (V.fromList ps) gen -samplePoisson -  :: Double   -- ^ Rate λ+samplePoisson+  :: Double   -- ^ rate λ   -> (MWC.GenIO -> IO Int) samplePoisson λ = \gen -> MWC.Probability.sample (MWC.Probability.poisson λ) gen -sampleDirichlet -  :: [Double] -- ^ Concentrations+sampleDirichlet+  :: [Double] -- ^ concentrations   -> (MWC.GenIO -> IO [Double]) sampleDirichlet xs = \gen -> MWC.Dist.dirichlet xs gen
src/Trace.hs view
@@ -1,14 +1,13 @@-{-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE TypeApplications #-} {-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE GADTs #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeOperators #-}-{-# LANGUAGE AllowAmbiguousTypes #-}-{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE UndecidableInstances #-} -{- | For recording samples and log-probabilities during model execution +{- | For recording samples and log-probabilities during model execution. -}  module Trace (@@ -31,17 +30,19 @@ import qualified Data.Map as Map import qualified OpenSum --- | The type of sample traces, mapping addresses of sample/observe operations to their primitive distributions and sampled values+{- | The type of sample traces, mapping addresses of sample/observe operations+     to their primitive distributions and sampled values.+-} type STrace = Map Addr (ErasedPrimDist, OpenSum PrimVal) --- | For converting sample traces, as used by simulation and inference, to output model environments-class FromSTrace env where +-- | For converting sample traces to model environments+class FromSTrace env where   -- | Convert a sample trace to a model environment   fromSTrace :: STrace -> Env env  instance FromSTrace '[] where   fromSTrace _ = nil- + instance (UniqueKey x env ~ 'True, KnownSymbol x, Eq a, OpenSum.Member a PrimVal, FromSTrace env) => FromSTrace ((x := a) : env) where   fromSTrace sMap = ECons (extractSamples (ObsVar @x, Proxy @a) sMap) (fromSTrace sMap) @@ -53,31 +54,33 @@  -- | Update a sample trace at an address updateSTrace :: (Show x, OpenSum.Member x PrimVal) =>-  -- | Address of sample site-     Addr       -  -- | Primitive distribution at address-  -> PrimDist x -  -- | Sampled value-  -> x -  -- | Previous sample trace-  -> STrace -  -- | Updated sample trace+  -- | address of sample site+     Addr+  -- | primitive distribution at address+  -> PrimDist x+  -- | sampled value+  -> x+  -- | previous sample trace   -> STrace+  -- | updated sample trace+  -> STrace updateSTrace α d x = Map.insert α (ErasedPrimDist d, OpenSum.inj x) --- | The type of log-probability traces, mapping addresses of sample/observe operations to their log probabilities+{- | The type of log-probability traces, mapping addresses of sample/observe operations+     to their log probabilities+-} type LPTrace = Map Addr Double  -- | Compute and update a log-probability trace at an address-updateLPTrace :: -  -- | Address of sample/observe site-     Addr -  -- | Primitive distribution at address-  -> PrimDist x -  -- | Sampled or observed value-  -> x -  -- | Prevous log-prob trace-  -> LPTrace -  -- | Updated log-prob trace+updateLPTrace ::+  -- | address of sample/observe site+     Addr+  -- | primitive distribution at address+  -> PrimDist x+  -- | sampled or observed value+  -> x+  -- | previous log-prob trace+  -> LPTrace+  -- | updated log-prob trace   -> LPTrace updateLPTrace α d x = Map.insert α (logProb d x)
src/Util.hs view
@@ -1,4 +1,4 @@-{- | Some small utility functions -}+{- | Some minor utility functions -}  module Util (     boolToInt