diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,4 @@ +# Change Log + +All notable changes to this project will be documented in this file. + diff --git a/LICENSE.md b/LICENSE.md new file mode 100644 --- /dev/null +++ b/LICENSE.md @@ -0,0 +1,30 @@ +BSD 3-Clause License + +Copyright (c) 2022, Minh Nguyen + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,91 @@ +## 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 interative 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! + +#### 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. + +#### 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. + +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 + type LogRegrEnv = + '[ "y" ':= Bool, -- ^ output + "m" ':= Double, -- ^ mean + "b" ':= Double -- ^ intercept + ] + + logRegr + :: (Observable env "y" Bool + , Observables env '["m", "b"] Double) + => [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 + let sigmoid x = 1.0 / (1.0 + exp((-1.0) * x)) + ys <- foldM (\ys x -> do + p <- normal' (m * x + b) sigma + y <- bernoulli (sigmoid p) #y + return (y:ys)) [] xs + return (reverse 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. + + 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. + + Below simulates from a logistic regression model using model parameters `m = 2` and `b = -0.15` but provides no values for `y`: this will result in `m` and `b` being *observed* but `y` being *sampled*. + ```haskell + simulateLogRegr :: Sampler [(Double, Bool)] + simulateLogRegr = do + -- | Specify the model inputs + let xs = map (/50) [(-50) .. 50] + -- | 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 + inferMHLogRegr :: Sampler [(Double, Double)] + inferMHLogRegr = do + -- | Simulate data from log regression + (xs, ys) <- unzip <$> simulateLogRegr + -- | 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"] + -- | Retrieve values sampled for #m and #b during MH + let m_samples = concatMap (get #m) mhTrace + 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. + + 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)] + ``` + + + diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ + +main = defaultMain diff --git a/examples/CoinFlip.hs b/examples/CoinFlip.hs new file mode 100644 --- /dev/null +++ b/examples/CoinFlip.hs @@ -0,0 +1,31 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE OverloadedLabels #-} +{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-} +{-# HLINT ignore "Redundant return" #-} +module CoinFlip where + +import Prog +import Effects.ObsReader +import Model +import PrimDist +import Effects.Dist +import Data.Kind (Constraint) +import Env + +-- ** Coin flip model (Section 5) +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 +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) + p <- call (Dist (UniformDist 0 1) maybe_p (Just "p")) + maybe_y <- call (Ask @env #y) + y <- call (Dist (BernoulliDist p) maybe_y (Just "p") ) + return y diff --git a/examples/DataSets.hs b/examples/DataSets.hs new file mode 100644 --- /dev/null +++ b/examples/DataSets.hs @@ -0,0 +1,20 @@ +module DataSets where + +n_counties :: Int +n_counties = 85 + +logRadon :: [Double] +logRadon = [0.8329091229351041, 0.8329091229351041, 1.0986122886681098, 0.0953101798043249, 1.1631508098056809, 0.9555114450274363, 0.4700036292457356, 0.0953101798043249, -0.2231435513142098, 0.262364264467491, 0.262364264467491, 0.336472236621213, 0.4054651081081644, -0.6931471805599453, 0.1823215567939548, 1.5260563034950492, 0.336472236621213, 0.7884573603642703, 1.791759469228055, 1.2237754316221157, 0.6418538861723948, 1.7047480922384253, 1.856297990365626, 0.6931471805599453, 1.9021075263969205, 1.1631508098056809, 1.9315214116032136, 1.9600947840472696, 2.0541237336955462, 1.667706820558076, 1.5260563034950492, 1.5040773967762742, 1.0647107369924282, 2.1041341542702074, 0.5306282510621705, 1.4586150226995167, 1.7047480922384253, 1.410986973710262, 0.8754687373538999, 1.0986122886681098, 0.4054651081081644, 1.2237754316221157, 1.0986122886681098, 0.6418538861723948, -1.203972804325936, 0.9162907318741552, 0.1823215567939548, 0.8329091229351041, -0.3566749439387324, 0.5877866649021191, 1.0986122886681098, 0.8329091229351041, 0.5877866649021191, 0.4054651081081644, 0.6931471805599453, 0.6418538861723948, 0.262364264467491, 1.4816045409242151, 1.5260563034950492, 1.856297990365626, 1.5475625087160128, 1.7578579175523736, 0.8329091229351041, -0.6931471805599453, 1.5475625087160128, 1.5040773967762742, 1.9021075263969205, 1.0296194171811583, 1.0986122886681098, 1.0986122886681098, 1.9878743481543453, 1.62924053973028, 0.9932517730102834, 1.62924053973028, 2.5726122302071057, 1.9878743481543453, 1.9315214116032136, 2.5572273113676265, 1.7749523509116736, 2.26176309847379, 1.8082887711792652, 1.3609765531356006, 2.667228206581955, 0.6418538861723948, 1.9459101490553128, 1.5686159179138452, 2.26176309847379, 0.9555114450274363, 1.916922612182061, 1.410986973710262, 2.322387720290225, 0.8329091229351041, 0.6418538861723948, 1.252762968495368, 1.7404661748405044, 1.4816045409242151, 1.3862943611198906, 0.336472236621213, 1.4586150226995167, -0.1053605156578262, 0.7419373447293773, 0.5306282510621705, 2.5649493574615367, 2.694627180770069, 1.5686159179138452, 2.272125885509337, -2.302585092994045, 1.3350010667323402, 2.0149030205422647, 0.6931471805599453, 1.6863989535702286, 1.410986973710262, 2.0541237336955462, 0.4054651081081644, 2.312535423847214, 2.2512917986064958, -0.1053605156578262, 1.5040773967762742, 1.62924053973028, 0.7884573603642703, 0.5877866649021191, 2.1041341542702074, 0.0, 2.5649493574615367, 0.9932517730102834, 1.2809338454620642, 3.284663565406204, 0.4700036292457356, 2.5726122302071057, 2.186051276738094, 2.975529566236472, 0.9555114450274363, 2.2082744135228043, 2.580216829592325, 1.308332819650179, 1.9459101490553128, 1.5892352051165808, 1.252762968495368, 0.0, 1.252762968495368, 1.0296194171811583, 0.4054651081081644, 1.9315214116032136, 2.4159137783010487, -2.302585092994045, 0.9555114450274363, 0.6418538861723948, 0.5306282510621705, 0.0953101798043249, 0.0, 1.0986122886681098, 1.5040773967762742, 0.4700036292457356, 1.4350845252893225, 0.9555114450274363, 1.916922612182061, 1.4816045409242151, 1.7227665977411035, 1.308332819650179, 1.0647107369924282, 2.6878474937846906, 1.916922612182061, 2.091864061678393, 0.9932517730102834, 1.0647107369924282, 1.5040773967762742, 0.5877866649021191, 0.7419373447293773, 0.7419373447293773, 0.4700036292457356, 2.272125885509337, 2.1041341542702074, 1.2809338454620642, -0.1053605156578262, 1.6486586255873816, 1.1939224684724346, 2.388762789235098, 2.116255514802552, 1.856297990365626, 1.5892352051165808, 1.8082887711792652, 0.1823215567939548, 2.1747517214841605, 2.186051276738094, 1.9315214116032136, 0.8754687373538999, 0.5306282510621705, 1.0647107369924282, 1.88706964903238, 0.5877866649021191, 1.5475625087160128, 1.2237754316221157, 1.5040773967762742, 3.05870707271538, 2.2192034840549946, 0.0, 1.6094379124341005, 1.62924053973028, 0.1823215567939548, 2.0412203288596382, 1.7047480922384253, 1.308332819650179, 1.6094379124341005, 1.5686159179138452, 0.4054651081081644, 1.252762968495368, 1.4586150226995167, 0.9555114450274363, 0.4054651081081644, 0.4054651081081644, 0.6931471805599453, 1.5892352051165808, 0.4054651081081644, 1.3609765531356006, 2.186051276738094, 1.4816045409242151, 1.5040773967762742, 1.5260563034950492, 0.8329091229351041, -0.5108256237659907, 1.7749523509116736, 1.7047480922384253, 1.9878743481543453, 1.7578579175523736, 2.0149030205422647, 1.5892352051165808, 1.9315214116032136, 1.8718021769015916, 1.3350010667323402, 1.7227665977411035, 2.066862759472976, 1.5040773967762742, 1.0296194171811583, 1.252762968495368, 1.4586150226995167, 0.8754687373538999, 0.336472236621213, 1.667706820558076, -1.6094379124341005, 0.9555114450274363, 1.1939224684724346, 1.1939224684724346, 2.272125885509337, 1.4586150226995167, 2.2082744135228043, 1.856297990365626, 3.487375077903208, 2.587764035227708, 0.8329091229351041, 1.7404661748405044, 2.667228206581955, 1.9459101490553128, 2.0412203288596382, 2.2925347571405443, 0.9932517730102834, 3.775057150354989, 1.6094379124341005, 1.6094379124341005, 1.2809338454620642, 1.5892352051165808, 1.7404661748405044, 1.2809338454620642, 1.3862943611198906, 1.916922612182061, 2.079441541679836, 1.2237754316221157, 0.7884573603642703, 0.5306282510621705, 1.410986973710262, 0.6418538861723948, 0.9555114450274363, 2.424802725718295, 0.9932517730102834, 1.3862943611198906, 2.0149030205422647, 0.336472236621213, 0.0, -0.6931471805599453, 0.9555114450274363, 1.8082887711792652, 0.7419373447293773, 1.7047480922384253, 1.1314021114911006, 1.0986122886681098, 1.7227665977411035, 1.4350845252893225, 1.3862943611198906, 2.70805020110221, 1.9878743481543453, 0.8754687373538999, 1.0647107369924282, 1.5040773967762742, 0.4700036292457356, 2.163323025660538, 1.7404661748405044, 2.163323025660538, 1.3609765531356006, 0.6418538861723948, 0.6931471805599453, 1.7227665977411035, 0.9555114450274363, -0.1053605156578262, 0.7884573603642703, 1.0647107369924282, 1.3862943611198906, 1.4816045409242151, 1.5686159179138452, 1.0647107369924282, 1.4350845252893225, 0.5306282510621705, 1.4816045409242151, -0.2231435513142098, 1.7227665977411035, 1.2237754316221157, 1.7227665977411035, 0.9555114450274363, 1.0296194171811583, 2.1400661634962708, 1.2237754316221157, 1.1939224684724346, 2.163323025660538, 0.5877866649021191, 1.7578579175523736, 2.5726122302071057, 1.0296194171811583, 1.5686159179138452, 1.7404661748405044, 2.631888840136646, 2.0412203288596382, 1.7578579175523736, 1.5475625087160128, 2.0412203288596382, 0.9932517730102834, 1.5260563034950492, 1.791759469228055, 0.8329091229351041, 0.9162907318741552, 1.410986973710262, 1.5475625087160128, 1.5475625087160128, 2.3978952727983707, 2.0412203288596382, 1.1314021114911006, 0.4700036292457356, 0.5306282510621705, 2.8094026953624978, 1.1631508098056809, 1.6486586255873816, 1.6094379124341005, 1.8082887711792652, 0.0, 0.6418538861723948, 1.3862943611198906, 1.7404661748405044, -0.6931471805599453, 0.9932517730102834, 1.308332819650179, 1.840549633397487, 3.1654750481410856, 1.3862943611198906, 1.0986122886681098, 1.1314021114911006, 1.5686159179138452, 1.1314021114911006, 1.4586150226995167, 1.3609765531356006, 1.1314021114911006, 1.4816045409242151, 1.0986122886681098, 1.252762968495368, 2.151762203259462, 2.2082744135228043, 1.5892352051165808, 1.308332819650179, 0.8329091229351041, 1.0647107369924282, -0.1053605156578262, 0.4700036292457356, 1.5475625087160128, 1.3350010667323402, 1.308332819650179, 1.1314021114911006, 0.8329091229351041, 0.6931471805599453, 0.9932517730102834, 0.6418538861723948, 0.9162907318741552, 1.4816045409242151, 0.9932517730102834, 0.1823215567939548, 1.2237754316221157, 0.9555114450274363, 2.2512917986064958, 0.336472236621213, 2.1400661634962708, 1.62924053973028, 1.0986122886681098, 2.580216829592325, 2.734367509419584, 0.6418538861723948, 1.3609765531356006, 2.079441541679836, 0.9932517730102834, 2.4336133554004498, 1.4350845252893225, 2.517696472610991, 1.916922612182061, 1.9459101490553128, 1.5260563034950492, 0.0, 0.5877866649021191, 0.4054651081081644, 0.7419373447293773, 0.0953101798043249, 0.0953101798043249, 1.0647107369924282, 0.336472236621213, 2.4336133554004498, 2.778819271990417, 0.336472236621213, 0.336472236621213, 0.5306282510621705, 0.0, 1.0647107369924282, -0.5108256237659907, 0.4700036292457356, 1.9740810260220096, -0.5108256237659907, 2.322387720290225, 1.4816045409242151, 1.2237754316221157, 1.0986122886681098, 2.533696813957432, 1.4586150226995167, 1.5260563034950492, 1.3862943611198906, 1.2237754316221157, 2.8678989020441064, 2.3702437414678603, 2.079441541679836, 1.2809338454620642, 1.88706964903238, 1.9459101490553128, 1.6486586255873816, 2.4932054526026954, 1.6486586255873816, 2.19722457733622, 1.7749523509116736, 1.5475625087160128, 1.3862943611198906, 0.4700036292457356, 3.173878458937465, 0.0, 0.4054651081081644, 0.1823215567939548, 1.0647107369924282, 3.8774315606585272, 0.0, 2.128231705849268, 1.4350845252893225, -0.5108256237659907, 1.916922612182061, 2.028148247292285, 2.23001440015921, -0.5108256237659907, 0.4700036292457356, 2.341805806147327, 1.3862943611198906, 0.6418538861723948, 2.302585092994046, 0.8754687373538999, 1.5040773967762742, 1.0647107369924282, 0.1823215567939548, 0.262364264467491, 0.5306282510621705, 3.2386784521643803, -2.302585092994045, 2.3702437414678603, 0.8754687373538999, 1.3862943611198906, 1.9878743481543453, 0.7884573603642703, 1.1939224684724346, -0.5108256237659907, 1.7578579175523736, 0.4054651081081644, 0.7884573603642703, 1.5040773967762742, 0.9162907318741552, 1.6094379124341005, 1.1314021114911006, 1.1314021114911006, 1.0647107369924282, 1.3862943611198906, 2.3978952727983707, 1.8718021769015916, 0.7419373447293773, 1.1314021114911006, 1.5260563034950492, 0.7884573603642703, 2.091864061678393, 0.336472236621213, 2.23001440015921, 0.1823215567939548, 2.3702437414678603, 3.1822118404966093, 2.2192034840549946, 2.501435951739211, 2.1041341542702074, 2.388762789235098, 1.4586150226995167, 2.760009940032921, 1.7047480922384253, 1.840549633397487, 2.282382385676526, 2.1041341542702074, 0.5306282510621705, 0.5306282510621705, 1.8718021769015916, 1.5040773967762742, 2.424802725718295, 2.312535423847214, 1.5260563034950492, 2.091864061678393, 0.8754687373538999, 1.1939224684724346, 1.62924053973028, 1.4350845252893225, 0.1823215567939548, 0.7419373447293773, 0.1823215567939548, 1.0986122886681098, 0.7884573603642703, 2.066862759472976, 1.3609765531356006, 0.9555114450274363, 1.0986122886681098, 0.5877866649021191, 0.9555114450274363, 2.2512917986064958, -0.3566749439387324, 1.0296194171811583, 0.1823215567939548, 0.7884573603642703, 2.4932054526026954, 2.5416019934645457, 1.1939224684724346, 1.4586150226995167, 1.3609765531356006, 1.3350010667323402, 1.7749523509116736, -0.916290731874155, 1.4350845252893225, 1.0647107369924282, 0.6931471805599453, 0.262364264467491, 0.262364264467491, 0.4700036292457356, 2.2512917986064958, 0.5877866649021191, 2.501435951739211, 1.4816045409242151, 1.9459101490553128, 0.4054651081081644, 0.9555114450274363, 2.272125885509337, 1.3609765531356006, 1.252762968495368, 1.9315214116032136, 1.308332819650179, 0.8329091229351041, 0.9932517730102834, 0.7884573603642703, 1.9600947840472696, 0.262364264467491, 1.3609765531356006, 1.2809338454620642, 1.4586150226995167, 0.5306282510621705, 1.0647107369924282, 2.163323025660538, 1.840549633397487, 1.667706820558076, 1.0296194171811583, 0.262364264467491, 1.2809338454620642, 1.7227665977411035, 2.322387720290225, 1.7227665977411035, 0.262364264467491, 1.6094379124341005, 1.410986973710262, 1.2809338454620642, 0.9555114450274363, 0.262364264467491, 1.0296194171811583, 0.5877866649021191, 1.1631508098056809, -0.2231435513142098, 0.0953101798043249, 0.6931471805599453, 1.3609765531356006, 2.19722457733622, 2.0149030205422647, 3.0349529867072724, 1.8082887711792652, 0.7884573603642703, 1.7749523509116736, 2.282382385676526, 1.8718021769015916, 1.5475625087160128, 1.7404661748405044, 2.9496883350525844, 0.9162907318741552, 1.1314021114911006, 1.6486586255873816, 2.0541237336955462, 2.1041341542702074, 1.5686159179138452, 2.1400661634962708, 0.5306282510621705, 1.8082887711792652, 0.1823215567939548, 2.4423470353692043, 1.4816045409242151, 1.308332819650179, 2.341805806147327, 1.252762968495368, 1.1631508098056809, 1.308332819650179, 1.0296194171811583, 1.410986973710262, 0.262364264467491, 0.5877866649021191, 1.4586150226995167, 2.9652730660692828, 2.2192034840549946, 0.7419373447293773, 2.4423470353692043, 2.33214389523559, 0.7884573603642703, 0.262364264467491, 1.1939224684724346, 0.7419373447293773, 1.4816045409242151, 0.8329091229351041, 1.7047480922384253, 3.2308043957334744, 1.6486586255873816, 0.8754687373538999, 1.1939224684724346, 0.9555114450274363, 1.0647107369924282, 1.1631508098056809, 0.5306282510621705, 1.5686159179138452, 1.410986973710262, 1.62924053973028, 0.4700036292457356, 1.5892352051165808, -0.1053605156578262, -0.5108256237659907, 0.9162907318741552, 0.8754687373538999, 1.5475625087160128, 2.4069451083182885, 2.70805020110221, 2.163323025660538, 1.5260563034950492, 0.4700036292457356, 1.3862943611198906, 0.6418538861723948, 0.5306282510621705, -0.5108256237659907, -0.6931471805599453, -0.5108256237659907, 2.1747517214841605, 0.5306282510621705, 0.4054651081081644, 2.1747517214841605, 2.4159137783010487, 0.4700036292457356, 0.1823215567939548, 0.0, -0.2231435513142098, 1.4586150226995167, 1.252762968495368, 0.7884573603642703, 1.0986122886681098, 0.6418538861723948, 0.6418538861723948, 0.9162907318741552, 0.5877866649021191, -0.1053605156578262, 2.468099531471619, 0.6418538861723948, 1.0647107369924282, 1.2809338454620642, 1.308332819650179, 1.2809338454620642, 1.1314021114911006, 1.1939224684724346, 1.1631508098056809, 1.2237754316221157, 0.5877866649021191, 1.7404661748405044, 1.252762968495368, 0.4700036292457356, 3.475067230228611, 0.1823215567939548, 0.7884573603642703, -0.1053605156578262, 0.4700036292457356, 0.336472236621213, 1.1631508098056809, 1.9878743481543453, 0.4054651081081644, 0.336472236621213, 0.4700036292457356, 1.62924053973028, 0.8754687373538999, 0.9162907318741552, 0.262364264467491, 1.7047480922384253, 0.1823215567939548, 0.4054651081081644, 1.9878743481543453, 0.1823215567939548, 1.2237754316221157, 1.1939224684724346, 0.4700036292457356, 1.308332819650179, -0.1053605156578262, 0.5306282510621705, 0.4054651081081644, 1.0296194171811583, 1.2237754316221157, 0.0, -0.3566749439387324, 0.7419373447293773, 0.6931471805599453, 0.0, 1.7047480922384253, 0.4700036292457356, 1.1631508098056809, 0.6418538861723948, 0.0, 1.2237754316221157, 0.5877866649021191, 1.1631508098056809, -0.2231435513142098, 1.4816045409242151, 0.4054651081081644, 0.6418538861723948, 0.4700036292457356, 0.8329091229351041, 0.9162907318741552, 1.0296194171811583, 0.5877866649021191, 0.1823215567939548, 0.6418538861723948, -1.203972804325936, 0.8329091229351041, 1.5475625087160128, 0.7884573603642703, 0.7419373447293773, -0.2231435513142098, 1.8718021769015916, 1.1314021114911006, 0.7419373447293773, 0.0, 1.2237754316221157, 0.6418538861723948, 0.6418538861723948, 0.8329091229351041, 1.4816045409242151, 2.028148247292285, 1.8718021769015916, 2.128231705849268, 0.7884573603642703, 1.2237754316221157, 0.336472236621213, 1.62924053973028, 0.0953101798043249, 1.9600947840472696, 1.7578579175523736, 2.322387720290225, 1.9021075263969205, 0.9932517730102834, 1.2237754316221157, 0.4700036292457356, 1.62924053973028, 2.0149030205422647, 2.681021528714291, 0.6418538861723948, 2.0149030205422647, 0.9932517730102834, 1.3350010667323402, 0.6931471805599453, 0.8329091229351041, 1.62924053973028, 2.001480000210124, 1.3350010667323402, 1.0986122886681098, 1.5040773967762742, 2.1400661634962708, 1.6486586255873816, 1.308332819650179, 0.4700036292457356, 2.163323025660538, 2.3702437414678603, 2.091864061678393, 1.5260563034950492, 1.1314021114911006, 0.9162907318741552, 0.4700036292457356, 1.5892352051165808, 1.9315214116032136, 0.7884573603642703, 1.8082887711792652, 1.0986122886681098, 1.916922612182061, 2.9652730660692828, 1.410986973710262, 1.791759469228055, 2.2082744135228043, 2.1400661634962708, 0.1823215567939548, 1.1631508098056809, 2.451005098112319, 2.272125885509337, 1.0986122886681098, -0.2231435513142098, 1.1939224684724346, 1.5686159179138452, 1.5892352051165808, -0.6931471805599453, 2.2407096892759584, 0.5877866649021191, 0.0, 2.33214389523559, 2.0541237336955462, 0.8329091229351041, 1.88706964903238, 2.509599262378372, 1.5475625087160128, 1.840549633397487, 1.88706964903238, 1.0647107369924282, 0.6931471805599453, 0.262364264467491, 0.9162907318741552, 0.0953101798043249, 0.262364264467491, 0.5306282510621705, -0.1053605156578262, 0.5877866649021191, 1.5686159179138452, 0.5877866649021191, 1.2237754316221157, -0.1053605156578262, 2.2925347571405443, 1.6863989535702286, 2.151762203259462, 0.6931471805599453, 1.9021075263969205, 1.3609765531356006, 1.791759469228055, 1.6094379124341005, 0.9555114450274363, 2.379546134130174, 0.9162907318741552, 0.7884573603642703, 1.5686159179138452, 1.3350010667323402, 2.602689685444384, 1.0986122886681098, 1.4816045409242151, 1.3609765531356006, 0.6418538861723948, 0.4700036292457356, 0.6418538861723948, 0.336472236621213, 1.9021075263969205, 3.0204248861443626, 1.8082887711792652, 2.631888840136646, 2.33214389523559, 1.7578579175523736, 2.2407096892759584, 1.252762968495368, 1.4350845252893225, 2.4595888418037104, 1.9878743481543453, 1.5686159179138452, 0.6418538861723948, -0.2231435513142098, 1.5686159179138452, 2.33214389523559, 2.4336133554004498, 2.0412203288596382, 2.476538400117484, -0.5108256237659907, 1.916922612182061, 1.6863989535702286, 1.1631508098056809, 0.7884573603642703, 2.001480000210124, 1.6486586255873816, 0.8329091229351041, 0.8754687373538999, 2.772588722239781, 2.26176309847379, 1.8718021769015916, 1.5260563034950492, 1.62924053973028, 1.3350010667323402, 1.0986122886681098] + +countyNames :: [String] +countyNames = [" AITKIN"," ANOKA"," BECKER"," BELTRAMI"," BENTON"," BIG STONE"," BLUE EARTH"," BROWN"," CARLTON"," CARVER"," CASS"," CHIPPEWA"," CHISAGO"," CLAY"," CLEARWATER"," COOK"," COTTONWOOD"," CROW WING"," DAKOTA"," DODGE"," DOUGLAS"," FARIBAULT"," FILLMORE"," FREEBORN"," GOODHUE"," HENNEPIN"," HOUSTON"," HUBBARD"," ISANTI"," ITASCA"," JACKSON"," KANABEC"," KANDIYOHI"," KITTSON"," KOOCHICHING"," LAC QUI PARLE"," LAKE"," LAKE OF THE WOODS"," LE SUEUR"," LINCOLN"," LYON"," MAHNOMEN"," MARSHALL"," MARTIN"," MCLEOD"," MEEKER"," MILLE LACS"," MORRISON"," MOWER"," MURRAY"," NICOLLET"," NOBLES"," NORMAN"," OLMSTED"," OTTER TAIL"," PENNINGTON"," PINE"," PIPESTONE"," POLK"," POPE"," RAMSEY"," REDWOOD"," RENVILLE"," RICE"," ROCK"," ROSEAU"," SCOTT"," SHERBURNE"," SIBLEY"," ST LOUIS"," STEARNS"," STEELE"," STEVENS"," SWIFT"," TODD"," TRAVERSE"," WABASHA"," WADENA"," WASECA"," WASHINGTON"," WATONWAN"," WILKIN"," WINONA"," WRIGHT"," YELLOW MEDICINE" ] + +-- length = 919 +countyIdx :: [Int] +countyIdx = + [0,0,0,0,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,2,2,2,3,3,3,3,3,3,3,4,4,4,4,5,5,5,6,6,6,6,6,6,6,6,6,6,6,6,6,6,7,7,7,7,8,8,8,8,8,8,8,8,8,8,9,9,9,9,9,9,10,10,10,10,10,11,11,11,11,12,12,12,12,12,12,13,13,13,13,13,13,13,13,13,13,13,13,13,13,14,14,14,14,15,15,16,16,16,16,17,17,17,17,17,17,17,17,17,17,17,17,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,18,19,19,19,20,20,20,20,20,20,20,20,20,21,21,21,21,21,21,22,22,23,23,23,23,23,23,23,23,23,24,24,24,24,24,24,24,24,24,24,24,24,24,24,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,25,26,26,26,26,26,26,27,27,27,27,27,28,28,28,29,29,29,29,29,29,29,29,29,29,29,30,30,30,30,30,31,31,31,31,32,32,32,32,33,33,33,34,34,34,34,34,34,34,35,35,36,36,36,36,36,36,36,36,36,37,37,37,37,38,38,38,38,38,39,39,39,39,40,40,40,40,40,40,40,40,41,42,42,42,42,42,42,42,42,42,43,43,43,43,43,43,43,44,44,44,44,44,44,44,44,44,44,44,44,44,45,45,45,45,45,46,46,47,47,47,47,47,47,47,47,47,48,48,48,48,48,48,48,48,48,48,48,48,48,49,50,50,50,50,51,51,51,52,52,52,53,53,53,53,53,53,53,53,53,53,53,53,53,53,53,53,53,53,53,53,53,53,53,54,54,54,54,54,54,54,54,55,55,55,56,56,56,56,56,56,57,57,57,57,58,58,58,58,59,59,60,60,60,60,60,60,60,60,60,60,60,60,60,60,60,60,60,60,60,60,60,60,60,60,60,60,60,60,60,60,60,60,61,61,61,61,61,62,62,62,63,63,63,63,63,63,63,63,63,63,63,64,64,65,65,65,65,65,65,65,65,65,65,65,65,65,65,66,66,66,66,66,66,66,66,66,66,66,66,66,67,67,67,67,67,67,67,67,68,68,68,68,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,69,70,70,70,70,70,70,70,70,70,70,70,70,70,70,70,70,70,70,70,70,70,70,70,70,70,71,71,71,71,71,71,71,71,71,71,72,72,73,73,73,73,74,74,74,75,75,75,75,76,76,76,76,76,76,76,77,77,77,77,77,78,78,78,78,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,79,80,80,80,81,82,82,82,82,82,82,82,82,82,82,82,82,82,83,83,83,83,83,83,83,83,83,83,83,83,83,84,84] + +-- length = 919 +dataFloorValues :: [Int] +dataFloorValues = + [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 1, 1, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 0, 1, 0, 0, 0, 0, 1, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 1, 1, 0, 1, 1, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 1, 1, 0, 1, 0, 0, 0, 1, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1, 1, 0, 1, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0] diff --git a/examples/HMM.hs b/examples/HMM.hs new file mode 100644 --- /dev/null +++ b/examples/HMM.hs @@ -0,0 +1,93 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE OverloadedLabels #-} +{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-} +{-# HLINT ignore "Redundant return" #-} + +module HMM where + +import Model +import Inference.SIM as SIM +import Inference.LW as LW +import Sampler +import Control.Monad +import Data.Kind (Constraint) +import Env +import Util + +-- ** HMM Loop +type HMMEnv = + '[ "trans_p" ':= Double, + "obs_p" ':= Double, + "y" ':= Int + ] + +hmmFor :: (Observable env "y" Int, Observables env '["obs_p", "trans_p"] Double) => + Int -> Int -> Model env ts Int +hmmFor n x = do + trans_p <- uniform 0 1 #trans_p + obs_p <- uniform 0 1 #obs_p + let hmmLoop i x_prev | i < n = do + dX <- boolToInt <$> bernoulli' trans_p + let x = x_prev + dX + binomial x obs_p #y + hmmLoop (i - 1) x + | otherwise = return x_prev + hmmLoop 0 x + +simulateHMM :: Sampler (Int, Env HMMEnv) +simulateHMM = do + let x_0 = 0; n = 10 + env = #trans_p := [0.5] <:> #obs_p := [0.8] <:> #y := [] <:> nil + SIM.simulate (hmmFor n) env 0 + +inferLwHMM :: Sampler [(Env HMMEnv, Double)] +inferLwHMM = do + let x_0 = 0; n = 10 + 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 +transModel :: Double -> Int -> Model env ts Int +transModel transition_p x_prev = do + dX <- boolToInt <$> bernoulli' transition_p + return (x_prev + dX) + +obsModel :: (Observable env "y" Int) + => Double -> Int -> Model env ts Int +obsModel observation_p x = do + y <- binomial x observation_p #y + return y + +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 + y_i <- obsModel observation_p x_i + return x_i + +hmm :: (Observable env "y" Int, Observables env '["obs_p", "trans_p"] Double) + => Int -> (Int -> Model env ts Int) +hmm n x = do + trans_p <- uniform 0 1 #trans_p + obs_p <- uniform 0 1 #obs_p + foldr (>=>) return (replicate n (hmmNode trans_p obs_p)) x + +-- ** 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 + +hmmGen :: Model env ts ps1 -> Model env ts ps2 + -> TransModel env ts ps1 lat -> ObsModel env ts ps2 lat obs + -> Int -> lat -> Model env ts lat +hmmGen transPrior obsPrior transModel obsModel n x_0 = do + ps1 <- transPrior + ps2 <- obsPrior + let hmmNode x = do + x' <- transModel ps1 x + y' <- obsModel ps2 x' + return x' + foldl (>=>) return (replicate n hmmNode) x_0 diff --git a/examples/LDA.hs b/examples/LDA.hs new file mode 100644 --- /dev/null +++ b/examples/LDA.hs @@ -0,0 +1,108 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE Trustworthy #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE Rank2Types #-} +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE OverloadedLabels #-} +{-# LANGUAGE QuantifiedConstraints #-} +{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-} +{-# HLINT ignore "Use camelCase" #-} + +module LDA where + +import Model +import Sampler +import Control.Monad +import Data.Kind (Constraint) +import Env +import Inference.SIM as SIM ( simulate ) +import Inference.MH as MH + + +-- ** Latent dirichlet allocation (topic model) +type TopicEnv = + '[ "θ" ':= [Double], + "φ" ':= [Double], + "w" ':= String + ] + +topicWordPrior :: Observable env "φ" [Double] + => [String] -> 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 +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) + => [String] + -> Int + -> Int + -> Model env ts [String] +topicModel vocab n_topics n_words = do + -- Generate distribution over words for each topic + topic_word_ps <- replicateM n_topics $ topicWordPrior vocab + -- Generate distribution over topics for a given document + doc_topic_ps <- docTopicPrior n_topics + replicateM n_words (do z <- discrete' doc_topic_ps + let word_ps = topic_word_ps !! z + wordDist vocab word_ps) + +-- | 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 + -> Model env ts [[String]] +topicModels vocab n_topics doc_words = do + mapM (topicModel vocab n_topics) doc_words + +-- | Simulating from topic model +-- | Possible vocabulary +vocab :: [String] +vocab = ["DNA", "evolution", "parsing", "phonology"] + +simLDA :: Sampler [String] +simLDA = do + let n_words = 100 + 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 + return words + +-- | MH Inference from topic model +-- | Document of words to perform inference over +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"] + +mhLDA :: Sampler ([[Double]], [[Double]]) +mhLDA = do + -- Do MH inference over the topic model using the above data + let n_words = 100 + n_topics = 2 + env_mh_in = #θ := [] <:> #φ := [] <:> #w := topic_data <:> nil + env_mh_outs <- MH.mh 500 (topicModel vocab n_topics) (n_words, env_mh_in) ["φ", "θ"] + -- Draw the most recent sampled parameters from MH + let env_pred = head env_mh_outs + θs = get #θ env_pred + φs = get #φ env_pred + return (θs, φs) diff --git a/examples/LinRegr.hs b/examples/LinRegr.hs new file mode 100644 --- /dev/null +++ b/examples/LinRegr.hs @@ -0,0 +1,67 @@ + +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE OverloadedLabels #-} +{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-} +{-# HLINT ignore "Redundant return" #-} +{-# LANGUAGE MonoLocalBinds #-} + +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 Data.Kind (Constraint) +import Env +import Util + +-- ** (Section 1) Linear regression +linRegr :: forall env rs . Member Sample rs => + Observables env '["y", "m", "c", "σ"] Double => + Double -> Model env rs Double +linRegr x = do + 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 +simulateLinRegr :: Sampler [(Double, Double)] +simulateLinRegr = do + let xs = [0 .. 100] + env = (#m := [3.0]) <:> (#c := [0]) <:> (#σ := [1]) <:> (#y := []) <:> nil + 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 +inferLwLinRegr :: Sampler [(Double, Double)] +inferLwLinRegr = do + let xs = [0 .. 100] + xys = [(x, env) | x <- xs, let env = (#m := []) <:> (#c := []) <:> (#σ := []) <:> (#y := [3*x]) <:> nil] + lwTrace <- mapM (LW.lw 200 linRegr) xys + let -- Get output of LW and extract mu samples + (env_outs, ps) = unzip $ concat lwTrace + mus = concatMap (get #m) env_outs + return $ zip mus ps + +-- Perform Metropolis-Hastings inference over linear regression +inferMhLinRegr :: Sampler [Double] +inferMhLinRegr = do + let xs = [0 .. 100] + xys = [(x, env) | x <- xs, let env = (#m := []) <:> (#c := []) <:> (#σ := []) <:> (#y := [3*x]) <:> nil] + 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 + return mus diff --git a/examples/LogRegr.hs b/examples/LogRegr.hs new file mode 100644 --- /dev/null +++ b/examples/LogRegr.hs @@ -0,0 +1,83 @@ +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedLabels #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE MonoLocalBinds #-} + +module LogRegr where + +import Control.Monad ( foldM ) +import Model ( bernoulli, gamma', normal, normal', Model ) +import Env ( (<:>), nil, Assign((:=)), Env, Observable(get), Observables ) +import Sampler ( Sampler ) +import Inference.SIM as SIM ( simulate ) +import Inference.MH as MH ( mh ) +import Inference.LW as LW ( lw ) + +-- | Logistic regression model +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. + -- 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] +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 + -- 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) + +-- | SIM from logistic regression +simulateLogRegr :: Sampler [(Double, Bool)] +simulateLogRegr = do + -- First declare the model inputs + let xs = map (/50) [(-50) .. 50] + -- Define a model environment to simulate from, providing observed values for the model parameters + env = (#y := []) <:> (#m := [2]) <:> (#b := [-0.15]) <:> nil + -- Call simulate on logistic regression + (ys, envs) <- SIM.simulate logRegr env xs + return (zip xs ys) + +-- | Likelihood-weighting over logistic regression +inferLwLogRegr :: Sampler [(Double, Double)] +inferLwLogRegr = do + -- Get values from simulating log regr + (xs, ys) <- unzip <$> simulateLogRegr + -- Define environment for inference, providing observed values for the model outputs + let env = (#y := ys) <:> (#m := []) <:> (#b := []) <:> nil + -- Run LW inference for 20000 iterations + lwTrace :: [(Env LogRegrEnv, Double)] <- LW.lw 20000 logRegr (xs, env) + let -- Get output of LW, extract mu samples, and pair with likelihood-weighting ps + (env_outs, ps) = unzip lwTrace + mus = concatMap (get #m) env_outs + return $ zip mus ps + +-- | Metropolis-Hastings inference over logistic regression +inferMHLogRegr :: Sampler [(Double, Double)] +inferMHLogRegr = do + -- Get values from simulating log regr + (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. + 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 + b_samples = concatMap (get #b) mhTrace + return (zip m_samples b_samples) diff --git a/examples/Main.hs b/examples/Main.hs new file mode 100644 --- /dev/null +++ b/examples/Main.hs @@ -0,0 +1,45 @@ +{-# LANGUAGE AllowAmbiguousTypes, PolyKinds #-} +{-# LANGUAGE DataKinds #-} + +module Main where + +import LinRegr +import LogRegr +import SIR +import LDA +import Radon +import School +import Sampler +import System.Environment (getArgs) + +printThenWrite :: Show a => a -> IO () +printThenWrite a = print a >> writeFile "model-output.txt" (show a) + +availableCommands = "[simLinRegr, lwLinRegr, mhLinRegr, simSIR, simSIRS, simSIRSV, mhSIR, simLogRegr, lwLogRegr, mhLogRegr, simLDA, mhLDA, simRadon, mhRadon, mhSchool]" + +parseArgs :: String -> IO () +parseArgs cmd = case cmd of + "simLinRegr" -> sampleIO simulateLinRegr >>= printThenWrite + "lwLinRegr" -> sampleIO inferLwLinRegr >>= printThenWrite + "simSIR" -> sampleIO simulateSIR >>= printThenWrite + "simSIRS" -> sampleIO simulateSIRS >>= printThenWrite + "simSIRSV" -> sampleIO simulateSIRSV >>= printThenWrite + "mhSIR" -> sampleIO inferSIR >>= printThenWrite + + "mhLinRegr" -> sampleIO inferMhLinRegr >>= printThenWrite + "simLogRegr" -> sampleIO simulateLogRegr >>= printThenWrite + "lwLogRegr" -> sampleIO inferLwLogRegr >>= printThenWrite + "mhLogRegr" -> sampleIO inferMHLogRegr >>= printThenWrite + "simLDA" -> sampleIO simLDA >>= printThenWrite + "mhLDA" -> sampleIO mhLDA >>= printThenWrite + "simRadon" -> sampleIO simRadon >>= printThenWrite + "mhRadon" -> sampleIO mhRadon >>= printThenWrite + "mhSchool" -> sampleIO mhSchool >>= printThenWrite + _ -> putStrLn $ "unrecognised command: " ++ cmd ++ "\n" + ++ "available commands: " ++ availableCommands + +main :: IO () +main = do + args <- getArgs + case args of [] -> print $ "no arguments provided to ProbFX. Available arguments: " ++ availableCommands + (a:as) -> parseArgs a diff --git a/examples/Radon.hs b/examples/Radon.hs new file mode 100644 --- /dev/null +++ b/examples/Radon.hs @@ -0,0 +1,89 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE OverloadedLabels #-} +{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-} +{-# HLINT ignore "Redundant return" #-} + +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 + +-- | 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] + +radonPrior :: Observables env '["mu_a", "mu_b", "sigma_a", "sigma_b"] Double + => Model env es (Double, Double, Double, Double) +radonPrior = do + mu_a <- normal 0 10 #mu_a + sigma_a <- halfNormal 5 #sigma_a + mu_b <- normal 0 10 #mu_b + 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 +radonModel :: Observables env '["mu_a", "mu_b", "sigma_a", "sigma_b", "a", "b", "log_radon"] Double + => Int -> [Int] -> [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 + a <- replicateM n_counties (normal mu_a sigma_a #a) -- length = 85 + -- Gradient for each county + b <- replicateM n_counties (normal mu_b sigma_b #b) -- length = 85 + -- Model error + eps <- halfCauchy' 5 + let -- Get county intercept for each datapoint + a_county_idx = map (a !!) county_idx + -- Get county gradient for each datapoint + b_county_idx = map (b !!) county_idx + floor_values = map fromIntegral floor_x + -- Get radon estimate for each data point + radon_est = zipWith (+) a_county_idx (zipWith (*) b_county_idx floor_values) + -- Sample radon amount for each data point + 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 (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 + +simRadon :: Sampler ([Double], [Double]) +simRadon = do + let env_in = mkRecordHLR ([1.45], [-0.68], [0.3], [0.2], [], [], []) + (bs, env_out) <- SIM.simulate (radonModel n_counties dataFloorValues countyIdx) env_in () + 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 +mhRadonpost :: Sampler ([Double], [Double]) +mhRadonpost = do + let env_in = mkRecordHLR ([], [], [], [], [], [], logRadon) + env_outs <- MH.mh 2000 (radonModel n_counties dataFloorValues countyIdx) ((), env_in) ["mu_a", "mu_b", "sigma_a", "sigma_b"] + let mu_a = concatMap (get #mu_a) env_outs + mu_b = concatMap (get #mu_b) env_outs + return (mu_a, mu_b) + +-- Return predictive posterior for intercepts and gradients +mhRadon :: Sampler ([Double], [Double]) +mhRadon = do + let env_in = mkRecordHLR ([], [], [], [], [], [], logRadon) + env_outs <- MH.mh 1500 (radonModel n_counties dataFloorValues countyIdx) ((), env_in) ["mu_a", "mu_b", "sigma_a", "sigma_b"] + let env_pred = head env_outs + as = get #a env_pred + bs = get #b env_pred + liftS $ print as + return (as, bs) diff --git a/examples/SIR.hs b/examples/SIR.hs new file mode 100644 --- /dev/null +++ b/examples/SIR.hs @@ -0,0 +1,233 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators, TypeApplications #-} +{-# LANGUAGE OverloadedLabels #-} +{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-} +{-# HLINT ignore "Redundant return" #-} + +module SIR where + +import Prog +import Effects.Writer +import Model +import Inference.SIM as SIM +import Inference.MH as MH +import Sampler +import Env +import Control.Monad + +import HMM + +-- ** The SIR model + +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 + +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 transition prior +transPriorSIR :: Observables env '["β", "γ"] Double + => Model env ts TransParamsSIR +transPriorSIR = do + pBeta <- gamma 2 1 #β + pGamma <- gamma 1 (1/8) #γ + return (TransParamsSIR pBeta pGamma) + +-- | 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 + +hmmSIR' :: (Observables env '["𝜉"] Int , Observables env '[ "β" , "γ" , "ρ"] Double) => Int -> Popl -> Model env es (Popl, [Popl]) +hmmSIR' n = handleWriterM . hmmSIR n + +type SIRenv = '["β" := Double, "γ" := Double, "ρ" := Double, "𝜉" := Int] + +-- ** Simulating from SIR model: ([(s, i, r)], [𝜉]) +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 + let 𝜉s :: [Reported] = get #𝜉 sim_env_out + 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 + 𝜉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) ["β", "ρ"] + 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' + +-- | SIR transition 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) + +-- | 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 + +-- || (Section 3.2, Fig 4b) SIM 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} + ((_, sir_trace), sim_env_out) <- SIM.simulate (hmmSIRS 100) sim_env_in sir_0 + let 𝜉s :: [Reported] = get #𝜉 sim_env_out + sirs = map (\(Popl s i recov) -> (s, i, recov)) 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 transition 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) + +-- | 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 + +-- | 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 + 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 + let 𝜉s :: [Reported] = get #𝜉 sim_env_out + sirvs = map (\(PoplV s i recov v) -> (s, i, recov, v)) sirv_trace + return (sirvs, 𝜉s) diff --git a/examples/SIRModular.hs b/examples/SIRModular.hs new file mode 100644 --- /dev/null +++ b/examples/SIRModular.hs @@ -0,0 +1,152 @@ +{-# 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 diff --git a/examples/School.hs b/examples/School.hs new file mode 100644 --- /dev/null +++ b/examples/School.hs @@ -0,0 +1,46 @@ + +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE OverloadedLabels #-} + +module School where + +import Model +import Inference.MH as MH +import Sampler +import Control.Monad +import Data.Kind (Constraint) +import Env + +-- | Hierarchical School Model +type SchEnv = '[ + "mu" ':= Double, + "theta" ':= [Double], + "y" ':= Double + ] + +schoolModel :: (Observables env '["mu", "y"] Double, Observable env "theta" [Double]) + => Int -> [Double] -> 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 +mhSchool :: Sampler ([Double], [[Double]]) +mhSchool = do + let n_schools = 8 + ys = [28, 8, -3, 7, -1, 1, 18, 12] + sigmas = [15, 10, 16, 11, 9, 11, 10, 18] + env = #mu := [] <:> #theta := [] <:> #y := ys <:> ENil + env_mh_out <- MH.mh 10000 (schoolModel n_schools ) (sigmas, env) ["mu", "theta"] + let mus = concatMap (get #mu) env_mh_out + thetas = concatMap (get #theta) env_mh_out + return (mus, thetas) diff --git a/prob-fx.cabal b/prob-fx.cabal new file mode 100644 --- /dev/null +++ b/prob-fx.cabal @@ -0,0 +1,101 @@ +cabal-version: 3.0 +name: prob-fx +version: 0.1.0.0 +license: BSD-3-Clause +license-file: LICENSE.md +copyright: 2022 Minh Nguyen +stability: experimental +author: Minh Nguyen +maintainer: minhnguyen1995@googlemail.com +homepage: https://github.com/min-nguyen/prob-fx/tree/hackage +synopsis: A library for modular probabilistic modelling +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 + of effect handlers. + +category: Statistics +build-type: Simple +extra-source-files: + README.md + CHANGELOG.md + +tested-with: + GHC == 8.6.5 + GHC == 8.8.4 + GHC == 8.10.4 + GHC == 9.0.1 + +library + exposed-modules: + Effects.Dist, + Effects.Lift, + Effects.ObsReader, + Effects.State, + Effects.Writer, + Inference.LW, + Inference.MH, + Inference.SIM, + Env, + FindElem, + Model, + OpenSum, + PrimDist, + Prog, + Sampler, + Trace, + Util + + -- Modules included in this library but not exported. + -- other-modules: + + -- LANGUAGE extensions used by modules in this package. + -- other-extensions: + + -- Other library packages from which modules are imported. + build-depends: base >= 4.11 && <= 4.17, + ghc-prim >= 0.5.3 && < 0.8, + deepseq >= 1.4.4 && < 1.5, + containers >= 0.6.0 && < 0.7, + primitive >= 0.7.4 && < 0.8, + transformers >= 0.5.6 && < 0.6, + random >= 1.2.1 && < 1.3, + mtl >= 2.2.2 && < 2.3, + vector >= 0.12.3 && < 0.13, + dirichlet >= 0.1.0 && < 0.2, + log-domain >= 0.13.2 && < 0.14, + mwc-random >= 0.15.0 && < 0.16, + extensible >= 0.9 && < 0.10, + membership >= 0.0.1 && < 0.1, + lens >= 5.1.1 && < 5.2, + mwc-probability >= 2.3.1 && < 2.4, + statistics >= 0.16.1 && < 0.17, + criterion >= 1.5.13 && < 1.6, + split >= 0.2.3 && < 0.3 + + -- Directories containing source files. + hs-source-dirs: src + + -- Base language which the package is written in. + default-language: Haskell2010 + ghc-options: -funfolding-use-threshold=16 -fexcess-precision -optc-O3 -optc-ffast-math + +executable examples + default-language: Haskell2010 + hs-source-dirs: examples + main-is: Main.hs + other-modules: CoinFlip, + DataSets, + HMM, + LDA, + LinRegr, + LogRegr, + Radon, + School, + SIR, + SIRModular, + build-depends: base, + prob-fx, + lens, + extensible diff --git a/src/Effects/Dist.hs b/src/Effects/Dist.hs new file mode 100644 --- /dev/null +++ b/src/Effects/Dist.hs @@ -0,0 +1,89 @@ +{-# 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 #-} + +{- | The effect for primitive distributions +-} + +module Effects.Dist ( + -- ** Address + Tag + , Addr + -- ** Dist effect + , Dist(..) + , handleDist + -- ** Sample effect + , Sample(..) + -- ** Observe effect + , Observe(..) + ) where + +import Data.Map (Map) +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 +type Tag = String +-- | An observable variable name and the index of its run-time occurrence +type Addr = (Tag, Int) + +-- | The Dist effect +data Dist a = Dist + { getPrimDist :: PrimDist a -- ^ primitive distribution + , getObs :: Maybe a -- ^ optional observed value + , getTag :: Maybe Tag -- ^ optional observable variable name + } + +instance Show a => Show (Dist a) where + show (Dist d y tag) = "Dist(" ++ show d ++ ", " ++ show y ++ ", " ++ show tag ++ ")" + +instance Eq (Dist a) where + (==) (Dist d1 _ _) (Dist d2 _ _) = d1 == d2 + +-- | An effect for sampling from distirbutions +data Sample a where + Sample :: PrimDist a -- ^ Distribution to sample from + -> Addr -- ^ Address of @Sample@ operation + -> Sample a + +-- | An effect 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 a + +-- | Handle the @Dist@ effect to a @Sample@ or @Observe@ effect and assign address +handleDist :: (Member Sample es, Member Observe es) + => Prog (Dist : es) a -> Prog es a +handleDist = loop 0 Map.empty + where + loop :: (Member Sample es, Member Observe es) + => Int -> Map Tag Int -> Prog (Dist : es) a -> Prog es a + loop _ _ (Val x) = return x + loop counter tagMap (Op u k) = case discharge u of + Right (Dist d maybe_y maybe_tag) -> + case maybe_y of + Just y -> do call (Observe d y (tag, tagIdx)) >>= k' + Nothing -> do call (Sample d (tag, tagIdx)) >>= k' + where tag = fromMaybe (show counter) maybe_tag + tagIdx = Map.findWithDefault 0 tag tagMap + tagMap' = Map.insert tag (tagIdx + 1) tagMap + k' = loop (counter + 1) tagMap' . k + Left u' -> Op u' (loop counter tagMap . k) + diff --git a/src/Effects/Lift.hs b/src/Effects/Lift.hs new file mode 100644 --- /dev/null +++ b/src/Effects/Lift.hs @@ -0,0 +1,36 @@ +{-# 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 +-} + +module Effects.Lift ( + Lift(..) + , lift + , handleLift) where + +import Prog ( call, Member(prj), Prog(..) ) +import Data.Function (fix) + +-- | Lift effect +newtype Lift m a = Lift (m a) + +-- | Lift a monadic computation @m a@ into the effect @Lift m@ +lift :: (Member (Lift m) es) => m a -> Prog es a +lift = call . Lift + +-- | Handle @Lift m@ as the last effect +handleLift :: forall m w. Monad m => Prog '[Lift m] w -> m w +handleLift (Val x) = return x +handleLift (Op u q) = case prj u of + Just (Lift m) -> m >>= handleLift . q + Nothing -> error "Impossible: Nothing cannot occur" + diff --git a/src/Effects/ObsReader.hs b/src/Effects/ObsReader.hs new file mode 100644 --- /dev/null +++ b/src/Effects/ObsReader.hs @@ -0,0 +1,43 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeOperators, TypeApplications #-} +{-# LANGUAGE AllowAmbiguousTypes #-} + +{- | An effect for reading observable variables from a model environment -} + +module Effects.ObsReader ( + ObsReader(..) + , ask + , handleRead) where + +import Prog ( call, discharge, Member, Prog(..) ) +import Env ( Env, ObsVar, Observable(..) ) +import Util ( safeHead, safeTail ) + +-- | 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 + +-- | 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) +ask x = call (Ask @env x) + +-- | Handle the @Ask@ requests of observable variables +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 + Right (Ask x) -> + let vs = get x env + maybe_v = safeHead vs + env' = set x (safeTail vs) env + in handleRead env' (k maybe_v) + Left op' -> Op op' (handleRead env . k) diff --git a/src/Effects/State.hs b/src/Effects/State.hs new file mode 100644 --- /dev/null +++ b/src/Effects/State.hs @@ -0,0 +1,50 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeOperators #-} + +{- | The state effect +-} + +module Effects.State ( + State(..) + , get + , put + , modify + , handleState) where + +import Prog ( discharge, Member(inj), Prog(..) ) + +-- | The state effect +data State s a where + Get :: State s s + Put :: s -> State s () + +-- | Get the state +get :: (Member (State s) es) => Prog es s +get = Op (inj Get) Val + +-- | Set the state +put :: (Member (State s) es) => s -> Prog es () +put s = Op (inj $ Put s) Val + +-- | 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 + -> Prog es (a, s) +handleState s m = loop s m where + loop :: s -> Prog (State s ': es) a -> Prog es (a, s) + loop s (Val x) = return (x, s) + loop s (Op u k) = case discharge u of + Right Get -> loop s (k s) + Right (Put s') -> loop s' (k ()) + Left u' -> Op u' (loop s . k) diff --git a/src/Effects/Writer.hs b/src/Effects/Writer.hs new file mode 100644 --- /dev/null +++ b/src/Effects/Writer.hs @@ -0,0 +1,46 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeOperators #-} + +{- | Writer effect +-} + +module Effects.Writer ( + Writer(..) + , tell + , tellM + , handleWriter + , handleWriterM) where + +import Prog ( discharge, Member(inj), Prog(..) ) +import Model ( Model(..) ) + +-- | 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 () + +-- | Write to a stream @w@ +tell :: Member (Writer w) ts => w -> Prog ts () +tell w = Op (inj $ Tell w) Val + +-- | Write to a stream @w@ inside a @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) +handleWriter = loop mempty where + loop :: w -> Prog (Writer w ': ts) a -> Prog ts (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 m = Model $ handleWriter $ runModel m diff --git a/src/Env.hs b/src/Env.hs new file mode 100644 --- /dev/null +++ b/src/Env.hs @@ -0,0 +1,120 @@ +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilyDependencies #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE TypeApplications #-} +{-# 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. +-} + +module Env + ( -- * Observable variable + ObsVar(..) + , varToStr + -- * Model environment + , Assign(..) + , Env(..) + , (<:>) + , nil + , Observable(..) + , Observables(..) + , UniqueKey + , LookupType) where + +import Data.Kind ( Constraint ) +import Data.Proxy ( Proxy(Proxy) ) +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 +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 + +-- | 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) + +-- | Assign or associate a variable @x@ with a value of type @a@ +data Assign x a = x := a + +-- | Empty model environment +nil :: Env '[] +nil = ENil + +infixr 5 <:> +-- | Prepend a variable assignment to a model environment +(<:>) :: UniqueKey x env ~ True => Assign (ObsVar x) [a] -> Env env -> Env ((x ':= a) ': env) +(_ := as) <:> env = ECons as env + +instance (KnownSymbol x, Show a, Show (Env env)) => Show (Env ((x := a) ': env)) where + show (ECons a env) = varToStr (ObsVar @x) ++ ":=" ++ show a ++ ", " ++ show env +instance Show (Env '[]) where + show ENil = "[]" + +instance FindElem x ((x := a) : env) where + findElem = Idx 0 +instance {-# OVERLAPPABLE #-} FindElem x env => FindElem x ((x' := a) : env) where + findElem = Idx $ 1 + unIdx (findElem :: Idx x env) + +-- | Retrieve the type of an observable variable @x@ from an environment @env@ +type family LookupType x env where + LookupType x ((x := a) : env) = a + LookupType x ((x' := a) : env) = LookupType x env + +-- | Specifies that an environment @Env env@ has an observable variable @x@ whose observed values are of type @a@ +class (FindElem x env, LookupType x env ~ a) + => Observable env x a where + get :: ObsVar x -> Env env -> [a] + set :: ObsVar x -> [a] -> Env env -> Env env + +instance (FindElem x env, LookupType x env ~ a) + => Observable env x a where + get _ env = + let idx = unIdx $ findElem @x @env + f :: Int -> Env env' -> [a] + f n (ECons a env) = if n == 0 + then unsafeCoerce a + else f (n - 1) env + in f idx env + set _ a' env = + let idx = unIdx $ findElem @x @env + f :: Int -> Env env' -> Env env' + f n (ECons a env) = if n == 0 + then ECons (unsafeCoerce a') env + else ECons a (f (n - 1) env) + in f idx env + +-- | For each observable variable @x@ in @xs@, construct the constraint @Observable env x a@ +type family Observables env (ks :: [Symbol]) a :: Constraint where + Observables env (x ': xs) a = (Observable env x a, Observables env xs a) + Observables env '[] a = () + +-- | Check whether an observable variable @x@ is unique in model environment @env@ +type family UniqueKey x env where + UniqueKey x ((x ':= a) : env) = False + UniqueKey x ((x' ':= a) : env) = UniqueKey x env + UniqueKey x '[] = True diff --git a/src/FindElem.hs b/src/FindElem.hs new file mode 100644 --- /dev/null +++ b/src/FindElem.hs @@ -0,0 +1,36 @@ +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilyDependencies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE ScopedTypeVariables #-} + +{- | Auxiliary definitions for finding a type in a type-level list. +-} + +module FindElem ( + FindElem(..) + , Idx(..)) where + +import GHC.TypeLits ( TypeError, ErrorMessage(Text, (:<>:), (:$$:), ShowType) ) + +-- | Proof that @x@ is an element of the type-level list @xs@ +class FindElem x xs where + findElem :: Idx x xs + +-- | The integer index of @x@ in @xs@ +newtype Idx x xs = Idx {unIdx :: Int} + +instance FindElem x (x ': xs) where + findElem = Idx 0 + +instance {-# OVERLAPPABLE #-} FindElem x xs => FindElem x (x' : xs) where + findElem = Idx $ 1 + unIdx (findElem :: Idx x xs) + +instance TypeError ('Text "Cannot unify effect types." ':$$: + 'Text "Unhandled effect: " ':<>: 'ShowType x ':$$: + 'Text "Perhaps check the type of effectful computation and the sequence of handlers for concordance?") + => FindElem x '[] where + findElem = error "unreachable" diff --git a/src/Inference/LW.hs b/src/Inference/LW.hs new file mode 100644 --- /dev/null +++ b/src/Inference/LW.hs @@ -0,0 +1,59 @@ +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE TypeOperators #-} + +{- | Likelihood-Weighting inference +-} + +module Inference.LW ( + lw + , runLW + , handleObs) where + +import qualified Data.Map as Map +import Env ( Env ) +import Effects.ObsReader ( ObsReader ) +import Control.Monad ( replicateM ) +import Effects.Dist ( Dist, Observe(..), Sample ) +import Prog ( discharge, Member, Prog(..) ) +import PrimDist ( logProb ) +import Model ( handleCore, Model ) +import Sampler ( Sampler ) +import Effects.State ( modify, handleState, State ) +import Trace ( FromSTrace(..), STrace ) +import Inference.SIM (traceSamples, handleSamp) + +-- | 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)] +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 + -> 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) +handleObs logp (Val x) = return (x, exp logp) +handleObs logp (Op u k) = case discharge u of + Right (Observe d y α) -> do + let logp' = logProb d y + handleObs (logp + logp') (k y) + Left op' -> Op op' (handleObs logp . k) diff --git a/src/Inference/MH.hs b/src/Inference/MH.hs new file mode 100644 --- /dev/null +++ b/src/Inference/MH.hs @@ -0,0 +1,191 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE GADTs #-} +{-# OPTIONS_GHC -Wno-incomplete-patterns #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE ViewPatterns #-} + +{- | Metropolis-Hastings inference +-} + +module Inference.MH ( + mh + , mhStep + , runMH + , traceLPs + , handleSamp + , lookupSample + , accept) where + +import Control.Monad ( (>=>) ) +import Data.Kind (Type) +import Data.Map (Map) +import Data.Maybe ( fromJust ) +import Data.Set (Set, (\\)) +import Effects.Dist ( Addr, Tag, Observe(..), Sample(..), Dist ) +import Effects.ObsReader ( ObsReader ) +import Effects.State ( State, modify, handleState ) +import Env ( Env ) +import Inference.SIM (handleObs, traceSamples) +import Model ( Model, handleCore ) +import OpenSum (OpenSum(..)) +import PrimDist +import Prog ( Member(prj), EffectSum, Prog(..), discharge ) +import qualified Data.Map as Map +import qualified Data.Set as Set +import qualified OpenSum +import Sampler ( Sampler, liftS ) +import Trace ( LPTrace, FromSTrace(..), STrace, updateLPTrace ) +import Unsafe.Coerce ( unsafeCoerce ) + +-- | 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] +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) + -- Perform n MH iterations + mhTrace <- foldl (>=>) return (replicate n (mhStep env_0 (model x_0) tags)) [y0] + -- Return sample trace + return $ map (\((_, strace), _) -> fromSTrace strace) mhTrace + +-- | 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 + -> Sampler [((a, STrace), LPTrace)] +mhStep env model tags trace = do + let -- Get previous mh output + ((x, samples), logps) = head trace + + sampleSites = if null tags then samples + else Map.filterWithKey (\(tag, i) _ -> tag `elem` tags) samples + + α_samp_ind <- sample $ DiscrUniformDist 0 (Map.size sampleSites - 1) + let (α_samp, _) = Map.elemAt α_samp_ind sampleSites + + ((x', samples'), logps') <- runMH env samples α_samp model + + 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 env es a + -- | (Outpiut, sample trace, log) + -> Sampler ((a, STrace), LPTrace) +runMH env strace α_samp = + handleSamp strace α_samp . handleObs + . handleState Map.empty . handleState Map.empty + . traceLPs . traceSamples . handleCore env + +pattern Samp :: Member Sample es => PrimDist x -> Addr -> EffectSum es x +pattern Samp d α <- (prj -> Just (Sample d α)) + +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 +traceLPs (Val x) = return x +traceLPs (Op op k) = case op of + Samp (PrimDistPrf d) α -> + Op op (\x -> modify (updateLPTrace α d x) >> + traceLPs (k x)) + Obs d y α -> + Op op (\ x -> modify (updateLPTrace α d y) >> + 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 + -> Sampler a +handleSamp strace α_samp (Op op k) = case discharge op of + Right (Sample (PrimDistPrf d) α) -> + do x <- lookupSample strace d α α_samp + handleSamp strace α_samp (k x) + _ -> error "Impossible: Nothing cannot occur" +handleSamp _ _ (Val x) = return x + +-- | 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 + -> Sampler a +lookupSample samples d α α_samp + | α == α_samp = sample d + | otherwise = + case Map.lookup α samples of + Just (ErasedPrimDist d', x) -> do + if d == unsafeCoerce d' + then return (fromJust $ OpenSum.prj x) + else sample d + Nothing -> sample d + +-- | Compute acceptance probability +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 + -> IO Double +accept x0 _Ⲭ _Ⲭ' logℙ logℙ' = do + let _X'sampled = Set.singleton x0 `Set.union` (Map.keysSet _Ⲭ' \\ Map.keysSet _Ⲭ) + _Xsampled = Set.singleton x0 `Set.union` (Map.keysSet _Ⲭ \\ Map.keysSet _Ⲭ') + let dom_logα = log (fromIntegral $ Map.size _Ⲭ) - log (fromIntegral $ Map.size _Ⲭ') + let _Xlogα = foldl (\logα v -> logα + fromJust (Map.lookup v logℙ)) + 0 (Map.keysSet logℙ \\ _Xsampled) + let _X'logα = foldl (\logα v -> logα + fromJust (Map.lookup v logℙ')) + 0 (Map.keysSet logℙ' \\ _X'sampled) + return $ exp (dom_logα + _X'logα - _Xlogα) diff --git a/src/Inference/SIM.hs b/src/Inference/SIM.hs new file mode 100644 --- /dev/null +++ b/src/Inference/SIM.hs @@ -0,0 +1,78 @@ +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE FlexibleContexts #-} +{-# OPTIONS_GHC -Wno-incomplete-patterns #-} + +{- | Simulation +-} + +module Inference.SIM ( + simulate + , runSimulate + , traceSamples + , handleSamp + , handleObs) where + +import Data.Map (Map) +import Effects.Dist ( Observe(..), Sample(..), Dist ) +import Effects.ObsReader ( ObsReader ) +import Effects.State ( State, modify, handleState ) +import Env ( Env ) +import Model ( Model, handleCore ) +import OpenSum (OpenSum) +import PrimDist +import Prog ( Member(prj), Prog(..), discharge ) +import qualified Data.Map as Map +import qualified OpenSum +import Sampler ( Sampler ) +import Trace ( FromSTrace(..), STrace, updateSTrace ) +import Unsafe.Coerce (unsafeCoerce) + +-- | 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) +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) +runSimulate env + = handleSamp . handleObs . handleState Map.empty . traceSamples . handleCore env + +-- | 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 + Just (Sample (PrimDistPrf d) α) -> + Op op (\x -> do modify (updateSTrace α d x); + traceSamples (k x)) + Nothing -> Op op (traceSamples . k) + +-- | Handler @Observe@ operations by simply passing forward their observed value, performing no side-effects +handleObs :: Prog (Observe : es) a -> Prog es a +handleObs (Val x) = return x +handleObs (Op op k) = case discharge op of + Right (Observe d y α) -> handleObs (k y) + Left op' -> Op op' (handleObs . k) + +-- | Handle @Sample@ operations by using the @Sampler@ monad to draw from primitive distributions +handleSamp :: Prog '[Sample] a -> Sampler a +handleSamp (Val x) = return x +handleSamp (Op op k) = case discharge op of + Right (Sample (PrimDistPrf d) α) -> + do x <- sample d + handleSamp (k x) + _ -> error "Impossible: Nothing cannot occur" diff --git a/src/Model.hs b/src/Model.hs new file mode 100644 --- /dev/null +++ b/src/Model.hs @@ -0,0 +1,357 @@ +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MonoLocalBinds #-} + +{- | An algebraic effect embedding of probabilistic models. +-} + +module Model ( + Model(..) + , handleCore + -- * Distribution smart constructors + -- $Smart-Constructors + , bernoulli + , bernoulli' + , beta + , beta' + , binomial + , binomial' + , categorical + , categorical' + , cauchy + , cauchy' + , halfCauchy + , halfCauchy' + , deterministic + , deterministic' + , dirichlet + , dirichlet' + , discrete + , discrete' + , gamma + , gamma' + , normal + , normal' + , halfNormal + , halfNormal' + , poisson + , poisson' + , uniform + , uniform' + ) + where + +import Control.Monad ( ap ) +import Control.Monad.Trans.Class ( MonadTrans(lift) ) +import Effects.Dist ( handleDist, Dist(Dist), Observe, Sample ) +import Effects.Lift ( Lift(..) ) +import Effects.ObsReader ( ask, handleRead, ObsReader ) +import Effects.State ( State, modify, handleState ) +import Env ( varToStr, Env, ObsVar, Observable ) +import OpenSum (OpenSum) +import 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, + +2) an effect signature @es@ of the possible effects a model can invoke, and + +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@. +-} +newtype Model env es a = + Model { runModel :: (Member Dist es, Member (ObsReader env) es) => Prog es a } + deriving Functor + +instance Applicative (Model env es) where + pure x = Model $ pure x + (<*>) = ap + +instance Monad (Model env es) where + return = pure + Model f >>= x = Model $ do + f' <- f + runModel $ x f' + +{- | 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: + +@ +exampleModel :: Observable env "b" Bool => Model env es Bool +exampleModel = bernoulli 0.5 #b +@ + +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 + -> 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 + -> Model env es a +deterministic' x = Model $ do + call (Dist (DeterministicDist x) Nothing Nothing) + +dirichlet :: forall env es x. Observable env x [Double] => + [Double] + -> ObsVar x + -> Model env es [Double] +dirichlet xs field = Model $ do + let tag = Just $ varToStr field + maybe_y <- ask @env field + call (Dist (DirichletDist xs) maybe_y tag) + +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] + -> ObsVar x + -> Model env es Int +discrete ps field = Model $ do + let tag = Just $ varToStr field + 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@ + -> 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)] + -> ObsVar x + -> Model env es a +categorical xs field = Model $ do + let tag = Just $ varToStr field + 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)] + -> Model env es a +categorical' xs = Model $ do + call (Dist (CategoricalDist xs) Nothing Nothing) + +normal :: forall env es x. Observable env x Double => + Double + -> Double + -> ObsVar x + -> Model env es Double +normal mu sigma field = Model $ do + let tag = Just $ varToStr field + maybe_y <- ask @env field + call (Dist (Normal mu sigma) maybe_y tag) + +normal' :: + -- | Mean + Double + -- | Standard deviation + -> Double + -> Model env es Double +normal' mu sigma = Model $ do + call (Dist (Normal mu sigma) Nothing Nothing) + +halfNormal :: forall env es x. Observable env x Double => + Double + -> ObsVar x + -> Model env es Double +halfNormal sigma field = Model $ do + let tag = Just $ varToStr field + maybe_y <- ask @env field + call (Dist (HalfNormalDist sigma) maybe_y tag) + +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 + -> ObsVar x + -> Model env es Double +cauchy mu sigma field = Model $ do + let tag = Just $ varToStr field + maybe_y <- ask @env field + call (Dist (CauchyDist mu sigma) maybe_y tag) + +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 + -> ObsVar x + -> Model env es Double +halfCauchy sigma field = Model $ do + let tag = Just $ varToStr field + maybe_y <- ask @env field + call (Dist (HalfCauchyDist sigma) maybe_y tag) + +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 + -> ObsVar x + -> Model env es Bool +bernoulli p field = Model $ do + let tag = Just $ varToStr field + maybe_y <- ask @env field + call (Dist (BernoulliDist p) maybe_y tag) + +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 + -> ObsVar x + -> Model env es Double +beta α β field = Model $ do + let tag = Just $ varToStr field + maybe_y <- ask @env field + call (Dist (BetaDist α β) maybe_y tag) + +beta' :: + -- | Shape 1 (α) + 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 + -> ObsVar x + -> Model env es Int +binomial n p field = Model $ do + let tag = Just $ varToStr field + 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 + -> 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 + -> ObsVar x + -> Model env es Double +gamma k θ field = Model $ do + let tag = Just $ varToStr field + maybe_y <- ask @env field + call (Dist (GammaDist k θ) maybe_y tag) + +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 + -> ObsVar x + -> Model env es Double +uniform min max field = Model $ do + let tag = Just $ varToStr field + maybe_y <- ask @env field + call (Dist (UniformDist min max) maybe_y tag) + +uniform' :: + -- | 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 + -> ObsVar x + -> Model env es Int +poisson λ field = Model $ do + let tag = Just $ varToStr field + maybe_y <- ask @env field + call (Dist (PoissonDist λ) maybe_y tag) + +poisson' :: + -- | Rate (λ) + Double + -- | Number of events + -> Model env es Int +poisson' λ = Model $ do + call (Dist (PoissonDist λ) Nothing Nothing) + diff --git a/src/OpenSum.hs b/src/OpenSum.hs new file mode 100644 --- /dev/null +++ b/src/OpenSum.hs @@ -0,0 +1,64 @@ +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE EmptyCase #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilyDependencies #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE UndecidableInstances #-} + +{- | An open sum implementation for value types. +-} +module OpenSum ( + OpenSum(..) + , Member(..)) where + +import Data.Kind (Type, Constraint) +import Data.Typeable ( Typeable ) +import FindElem ( Idx(..), FindElem(..) ) +import GHC.TypeLits (Nat, KnownNat, natVal, TypeError, ErrorMessage (Text, (:$$:), (:<>:), ShowType)) +import qualified GHC.TypeLits as TL +import Unsafe.Coerce ( unsafeCoerce ) + +-- | Open sum of value types +data OpenSum (as :: [*]) where + UnsafeOpenSum :: Int -> a -> OpenSum as + +instance Eq (OpenSum '[]) where + x == _ = case x of {} + +instance forall a as. (Eq a, Eq (OpenSum as)) => Eq (OpenSum (a : as)) where + UnsafeOpenSum i _ == UnsafeOpenSum j _ | i /= j = False + UnsafeOpenSum 0 x == UnsafeOpenSum 0 y = + unsafeCoerce x == (unsafeCoerce y :: a) + UnsafeOpenSum i x == UnsafeOpenSum j y = + UnsafeOpenSum (i - 1) x == (UnsafeOpenSum (j - 1) y :: OpenSum as) + +instance forall a as. (Show a, Show (OpenSum as)) => Show (OpenSum (a : as)) where + show (UnsafeOpenSum i a) = + if i == 0 + then show (unsafeCoerce a :: a) + else show (UnsafeOpenSum (i - 1) a :: OpenSum as) + +instance {-# OVERLAPPING #-} Show a => Show (OpenSum '[a]) where + show (UnsafeOpenSum i a) = show (unsafeCoerce a :: a) + +-- | Safely inject and project a value into an open sum +class (FindElem a as) => Member (a :: *) (as :: [*]) where + inj :: a -> OpenSum as + prj :: OpenSum as -> Maybe a + +instance (Typeable a, a ~ a') => Member a '[a'] where + inj x = UnsafeOpenSum 0 x + prj (UnsafeOpenSum _ x) = Just (unsafeCoerce x) + +instance (FindElem a as) => Member a as where + inj = UnsafeOpenSum (unIdx (findElem :: Idx a as)) + prj = prj' (unIdx (findElem :: Idx a as)) + where prj' n (UnsafeOpenSum n' x) + | n == n' = Just (unsafeCoerce x) + | otherwise = Nothing diff --git a/src/PrimDist.hs b/src/PrimDist.hs new file mode 100644 --- /dev/null +++ b/src/PrimDist.hs @@ -0,0 +1,294 @@ +{-# 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 ScopedTypeVariables #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE EmptyCase #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE ImplicitParams #-} +{-# OPTIONS_GHC -Wno-incomplete-patterns #-} + +{- | A GADT encoding of (a selection of) primitive distributions + along with their corresponding sampling and density functions. +-} + +module PrimDist ( + -- * Primitive distribution + PrimDist(..) + , PrimVal + , IsPrimVal(..) + , pattern PrimDistPrf + , ErasedPrimDist(..) + -- * Sampling + , sample + -- * Density + , prob + , logProb) where + +import Data.Kind ( Constraint ) +import Data.Map (Map) +import Numeric.Log ( Log(Exp) ) +import OpenSum (OpenSum) +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 +import Statistics.Distribution ( ContDistr(density), DiscreteDistr(probability) ) +import Statistics.Distribution.Beta ( betaDistr ) +import Statistics.Distribution.Binomial ( binomial ) +import Statistics.Distribution.CauchyLorentz ( cauchyDistribution ) +import Statistics.Distribution.Dirichlet ( dirichletDensity, dirichletDistribution ) +import Statistics.Distribution.DiscreteUniform ( discreteUniformAB ) +import Statistics.Distribution.Gamma ( gammaDistr ) +import Statistics.Distribution.Normal ( normalDistr ) +import Statistics.Distribution.Poisson ( poisson ) +import Statistics.Distribution.Uniform ( uniformDistr ) +import Sampler +import Util ( boolToInt ) + +-- | Primitive distribution +data PrimDist a where + 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 + -> PrimDist Double + HalfCauchyDist + :: Double -- ^ Scale + -> PrimDist Double + DeterministicDist + :: (Eq a, Show a, OpenSum.Member a PrimVal) + => a -- ^ Value of probability @1@ + -> PrimDist a + 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 θ + -> PrimDist Double + Normal + :: Double -- ^ Mean + -> Double -- ^ Standard deviation + -> PrimDist Double + HalfNormalDist + :: Double -- ^ Standard deviation + -> PrimDist Double + PoissonDist + :: Double -- ^ Rate λ + -> PrimDist Int + UniformDist + :: Double -- ^ Lower-bound @a@ + -> Double -- ^ Upper-bound @b@ + -> PrimDist Double + +instance Eq (PrimDist a) where + (==) (Normal m s) (Normal m' s') = m == m' && s == s' + (==) (CauchyDist m s) (CauchyDist m' s') = m == m' && s == s' + (==) (HalfCauchyDist s) (HalfCauchyDist s') = s == s' + (==) (HalfNormalDist s) (HalfNormalDist s') = s == s' + (==) (BernoulliDist p) (BernoulliDist p') = p == p' + (==) (BinomialDist n p) (BinomialDist n' p') = n == n' && p == p' + (==) (DiscreteDist ps) (DiscreteDist ps') = ps == ps' + (==) (BetaDist a b) (BetaDist a' b') = a == a' && b == b' + (==) (GammaDist a b) (GammaDist a' b') = a == a' && b == b' + (==) (UniformDist a b) (UniformDist a' b') = a == a' && b == b' + (==) (DiscrUniformDist min max) (DiscrUniformDist min' max') = min == min' && max == max' + (==) (PoissonDist l) (PoissonDist l') = l == l' + (==) (CategoricalDist xs) (CategoricalDist xs') = xs == xs' + (==) (DirichletDist xs) (DirichletDist xs') = xs == xs' + (==) (DeterministicDist x) (DeterministicDist x') = x == x' + (==) _ _ = False + +instance Show a => Show (PrimDist a) where + show (CauchyDist mu sigma) = + "CauchyDist(" ++ show mu ++ ", " ++ show sigma ++ ", " ++ ")" + show (HalfCauchyDist sigma) = + "HalfCauchyDist(" ++ show sigma ++ ", " ++ ")" + show (Normal mu sigma) = + "Normal(" ++ show mu ++ ", " ++ show sigma ++ ", " ++ ")" + show (HalfNormalDist sigma) = + "HalfNormalDist(" ++ show sigma ++ ", " ++ ")" + show (BernoulliDist p) = + "BernoulliDist(" ++ show p ++ ", " ++ ")" + show (BinomialDist n p) = + "BinomialDist(" ++ show n ++ ", " ++ show p ++ ", " ++ ")" + show (DiscreteDist ps) = + "DiscreteDist(" ++ show ps ++ ", " ++ ")" + show (BetaDist a b) = + "BetaDist(" ++ show a ++ ", " ++ show b ++ "," ++ ")" + show (GammaDist a b) = + "GammaDist(" ++ show a ++ ", " ++ show b ++ "," ++ ")" + show (UniformDist a b) = + "UniformDist(" ++ show a ++ ", " ++ show b ++ "," ++ ")" + show (DiscrUniformDist min max) = + "DiscrUniformDist(" ++ show min ++ ", " ++ show max ++ ", " ++ ")" + show (PoissonDist l) = + "PoissonDist(" ++ show l ++ ", " ++ ")" + show (CategoricalDist xs) = + "CategoricalDist(" ++ show xs ++ ", " ++ ")" + show (DirichletDist xs) = + "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] + +-- | Proof that @x@ is a primitive value +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 +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 d = case d of + HalfCauchyDist {} -> IsPrimVal + CauchyDist {} -> IsPrimVal + Normal {} -> IsPrimVal + HalfNormalDist {} -> IsPrimVal + UniformDist {} -> IsPrimVal + DiscrUniformDist {} -> IsPrimVal + GammaDist {} -> IsPrimVal + BetaDist {} -> IsPrimVal + BinomialDist {} -> IsPrimVal + BernoulliDist {} -> IsPrimVal + CategoricalDist {} -> IsPrimVal + DiscreteDist {} -> IsPrimVal + PoissonDist {} -> IsPrimVal + DirichletDist {} -> IsPrimVal + DeterministicDist {} -> IsPrimVal + +-- | For erasing the types of primitive distributions +data ErasedPrimDist where + ErasedPrimDist :: forall a. Show a => PrimDist a -> ErasedPrimDist + +instance Show ErasedPrimDist where + show (ErasedPrimDist d) = show d + +-- | Draw a value from a primitive distribution in the @Sampler@ monad +sample :: + PrimDist a + -> Sampler a +sample (HalfCauchyDist σ ) = + createSampler (sampleCauchy 0 σ) >>= pure . abs +sample (CauchyDist μ σ ) = + createSampler (sampleCauchy μ σ) +sample (HalfNormalDist σ ) = + createSampler (sampleNormal 0 σ) >>= pure . abs +sample (Normal μ σ ) = + createSampler (sampleNormal μ σ) +sample (UniformDist min max ) = + createSampler (sampleUniform min max) +sample (DiscrUniformDist min max ) = + createSampler (sampleDiscreteUniform min max) +sample (GammaDist k θ ) = + createSampler (sampleGamma k θ) +sample (BetaDist α β ) = + createSampler (sampleBeta α β) +sample (BinomialDist n p ) = + createSampler (sampleBinomial n p) >>= pure . length . filter (== True) +sample (BernoulliDist p ) = + createSampler (sampleBernoulli p) +sample (CategoricalDist ps ) = + createSampler (sampleCategorical (V.fromList $ fmap snd ps)) >>= \i -> pure $ fst $ ps !! i +sample (DiscreteDist ps ) = + createSampler (sampleDiscrete ps) +sample (PoissonDist λ ) = + createSampler (samplePoisson λ) +sample (DirichletDist xs ) = + createSampler (sampleDirichlet xs) +sample (DeterministicDist x) = pure x + +-- | Compute the density of a primitive distribution generating an observed value +prob :: + -- | Distribution + PrimDist a + -- | Observed value + -> a + -- | Density + -> Double +prob (DirichletDist xs) ys = + let xs' = map (/(Prelude.sum xs)) xs + in if Prelude.sum xs' /= 1 then error "dirichlet can't normalize" else + case dirichletDistribution (UV.fromList xs') + of Left e -> error "dirichlet error" + Right d -> let Exp p = dirichletDensity d (UV.fromList ys) + in exp p +prob (HalfCauchyDist σ) y + = if y < 0 then 0 else + 2 * density (cauchyDistribution 0 σ) y +prob (CauchyDist μ σ) y + = density (cauchyDistribution μ σ) y +prob (HalfNormalDist σ) y + = if y < 0 then 0 else + 2 * density (normalDistr 0 σ) y +prob (Normal μ σ) y + = density (normalDistr μ σ) y +prob (UniformDist min max) y + = density (uniformDistr min max) y +prob (GammaDist k θ) y + = density (gammaDistr k θ) y +prob (BetaDist α β) y + = density (betaDistr α β) y +prob (DiscrUniformDist min max) y + = probability (discreteUniformAB min max) y +prob (BinomialDist n p) y + = probability (binomial n p) y +prob (BernoulliDist p) i + = probability (binomial 1 p) (boolToInt i) +prob d@(CategoricalDist ps) y + = 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 + +-- | Compute the log density of a primitive distribution generating an observed value +logProb :: + -- | Distribution + PrimDist a + -- | Observed value + -> a + -- | Log density + -> Double +logProb d = log . prob d diff --git a/src/Prog.hs b/src/Prog.hs new file mode 100644 --- /dev/null +++ b/src/Prog.hs @@ -0,0 +1,98 @@ +{-# LANGUAGE DataKinds #-} +{-# 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. +-} + +module Prog ( + -- * Effectful program + Prog(..) + , EffectSum + , Member(..) + -- * Auxiliary functions + , run + , call + , discharge) where + +import Control.Monad ( (>=>) ) +import Data.Kind (Constraint) +import FindElem ( Idx(unIdx), FindElem(..) ) +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. +data Prog es a where + 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 + -> Prog es a + +instance Functor (Prog es) where + fmap f (Val a) = Val (f a) + fmap f (Op fx k) = Op fx (fmap f . k) + +instance Applicative (Prog es) where + pure = Val + Val f <*> x = fmap f x + (Op fx k) <*> x = Op fx ((<*> x) . k) + +instance Monad (Prog es) where + return = Val + Val a >>= f = f a + Op fx k >>= f = Op fx (k >=> f) + +-- | An open sum for an effect signature @es@, containing an operation @e x@ where @e@ is in @es@ +data EffectSum (es :: [* -> *]) (x :: *) :: * where + EffectSum :: Int -> e x -> EffectSum es x + +-- | 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 + -- | Attempt to project an operation of type @e x@ out from an effect sum + prj :: EffectSum es x -> Maybe (e x) + +instance {-# INCOHERENT #-} (e ~ e') => Member e '[e'] where + inj x = EffectSum 0 x + prj (EffectSum _ x) = Just (unsafeCoerce x) + +instance (FindElem e es) => Member e es where + inj = EffectSum (unIdx (findElem :: Idx e es)) + prj = prj' (unIdx (findElem :: Idx e es)) + where prj' n (EffectSum n' x) + | n == n' = Just (unsafeCoerce x) + | otherwise = Nothing + +-- | Membership of many effects @es@ in @ess@ +type family Members (es :: [* -> *]) (ess :: [* -> *]) = (cs :: Constraint) | cs -> es where + Members (e ': es) ess = (Member e ess, Members es ess) + Members '[] ess = () + +-- | Run a pure computation +run :: Prog '[] a -> a +run (Val x) = x +run _ = error "'run' isn't defined for non-pure computations" + +-- | Call an operation of type @e x@ in a computation +call :: (Member e es) => e x -> Prog es x +call e = Op (inj e) Val + +-- | Discharges an effect @e@ from the front of an effect signature @es@ +discharge :: EffectSum (e ': es) x -> Either (EffectSum es x) (e x) +discharge (EffectSum 0 tv) = Right $ unsafeCoerce tv +discharge (EffectSum n rv) = Left $ EffectSum (n-1) rv + +-- | For pattern-matching against operations that belong in the tail of an effect signature +pattern Other :: EffectSum es x -> EffectSum (e ': es) x +pattern Other u <- (discharge -> Left u) diff --git a/src/Sampler.hs b/src/Sampler.hs new file mode 100644 --- /dev/null +++ b/src/Sampler.hs @@ -0,0 +1,145 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving #-} + +{- | An IO-based sampling monad +-} + +module Sampler ( + -- * Sampler monad + Sampler + , liftS + , sampleIO + , sampleIOFixed + , createSampler + -- * Sampling functions + -- $Sampling-functions + , sampleRandom + , sampleCauchy + , sampleNormal + , sampleUniform + , sampleDiscreteUniform + , sampleGamma + , sampleBeta + , sampleBernoulli + , sampleBinomial + , sampleCategorical + , sampleDiscrete + , samplePoisson + , sampleDirichlet + ) 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 +import qualified System.Random.MWC.Probability as MWC.Probability +import Statistics.Distribution ( ContGen(genContVar) ) +import Statistics.Distribution.CauchyLorentz ( cauchyDistribution ) +import System.Random.MWC ( initialize ) + +-- | Sampler type, for running IO computations alongside a random number generator +newtype Sampler a = Sampler {runSampler :: ReaderT MWC.GenIO IO a} + deriving (Functor, Applicative, Monad) + +-- | Lift an @IO@ computation into @Sampler@ +liftS :: IO a -> Sampler a +liftS f = Sampler $ lift f + +-- | Takes a @Sampler@, provides it a random generator, and runs the sampler in the @IO@ context +sampleIO :: Sampler a -> IO a +sampleIO m = MWC.createSystemRandom >>= (runReaderT . runSampler) m + +-- | Takes a @Sampler@, provides it a fixed generator, and runs the sampler in the @IO@ context +sampleIOFixed :: Sampler a -> IO a +sampleIOFixed m = MWC.create >>= (runReaderT . runSampler) m + +-- | Takes a @Sampler@, provides it a custom fixed generator, and runs the sampler in the @IO@ context +sampleIOCustom :: Int -> Sampler a -> IO a +sampleIOCustom n m = initialize (V.singleton (fromIntegral n :: Word32)) >>= (runReaderT . runSampler) m + +-- | Takes a distribution which awaits a generator, and returns a @Sampler@ +createSampler :: (MWC.GenIO -> IO a) -> Sampler a +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 +-} + +sampleRandom + :: MWC.GenIO + -> IO Double +sampleRandom = \gen -> MWC.uniform gen + +sampleCauchy + :: Double -- ^ Location + -> Double -- ^ Scale + -> (MWC.GenIO -> IO Double) +sampleCauchy μ σ = \gen -> genContVar (cauchyDistribution μ σ) gen + +sampleNormal + :: Double -- ^ Mean + -> Double -- ^ Standard deviation + -> (MWC.GenIO -> IO Double) +sampleNormal μ σ = \gen -> MWC.Dist.normal μ σ gen + +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 + -> (MWC.GenIO -> IO Int) +sampleDiscreteUniform min max = \gen -> MWC.uniformR (min, max) gen + +sampleGamma + :: Double -- ^ Shape k + -> Double -- ^ Scale θ + -> (MWC.GenIO -> IO Double) +sampleGamma k θ = \gen -> MWC.Dist.gamma k θ gen + +sampleBeta + :: Double -- ^ Shape α + -> Double -- ^ Shape β + -> (MWC.GenIO -> IO Double) +sampleBeta α β = \gen -> MWC.Dist.beta α β gen + +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 + -> (MWC.GenIO -> IO [Bool]) +sampleBinomial n p = \gen -> replicateM n (MWC.Dist.bernoulli p gen) + +sampleCategorical + :: V.Vector Double -- ^ Probabilities + -> (MWC.GenIO -> IO Int) +sampleCategorical ps = \gen -> MWC.Dist.categorical (ps) gen + +sampleDiscrete + :: [Double] -- ^ Probabilities + -> (MWC.GenIO -> IO Int) +sampleDiscrete ps = \gen -> MWC.Dist.categorical (V.fromList ps) gen + +samplePoisson + :: Double -- ^ Rate λ + -> (MWC.GenIO -> IO Int) +samplePoisson λ = \gen -> MWC.Probability.sample (MWC.Probability.poisson λ) gen + +sampleDirichlet + :: [Double] -- ^ Concentrations + -> (MWC.GenIO -> IO [Double]) +sampleDirichlet xs = \gen -> MWC.Dist.dirichlet xs gen diff --git a/src/Trace.hs b/src/Trace.hs new file mode 100644 --- /dev/null +++ b/src/Trace.hs @@ -0,0 +1,83 @@ +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE UndecidableInstances #-} + +{- | For recording samples and log-probabilities during model execution +-} + +module Trace ( + -- * Sample trace + STrace + , FromSTrace(..) + , updateSTrace + -- * Log-probability trace + , LPTrace + , updateLPTrace) where + +import Data.Map (Map) +import Data.Maybe ( fromJust ) +import Data.Proxy ( Proxy(..) ) +import Effects.Dist ( Addr ) +import PrimDist ( ErasedPrimDist(..), PrimVal, PrimDist, logProb ) +import Env ( UniqueKey, Assign((:=)), Env(ECons), ObsVar(..), varToStr, nil ) +import GHC.TypeLits ( KnownSymbol ) +import OpenSum (OpenSum) +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 +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 + -- | 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) + +extractSamples :: forall a x. (Eq a, OpenSum.Member a PrimVal) => (ObsVar x, Proxy a) -> STrace -> [a] +extractSamples (x, typ) = + map (fromJust . OpenSum.prj @a . snd . snd) + . Map.toList + . Map.filterWithKey (\(tag, idx) _ -> tag == varToStr x) + +-- | 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 + -> 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 +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 + -> LPTrace +updateLPTrace α d x = Map.insert α (logProb d x) diff --git a/src/Util.hs b/src/Util.hs new file mode 100644 --- /dev/null +++ b/src/Util.hs @@ -0,0 +1,30 @@ +{- | Some small utility functions -} + +module Util ( + boolToInt + , safeHead + , safeTail + , findIndexes) where + +-- | Return @True@ for @1@ and otherwise @False@ +boolToInt :: Bool -> Int +boolToInt True = 1 +boolToInt False = 0 + +-- | Safely attempt to return the head of a list +safeHead :: [a] -> Maybe a +safeHead [] = Nothing +safeHead (x:xs) = Just x + +-- | Return the tail of a list, behaving as the identity function upon an empty list +safeTail :: [a] -> [a] +safeTail [] = [] +safeTail (x:xs) = xs + +-- | Return all the positions that a value occurs within a list +findIndexes :: Eq a => [a] -> a -> [Int] +findIndexes xs a = reverse $ go xs 0 [] + where + go (x:xs) i inds = if x == a then go xs (i + 1) (i:inds) + else go xs (i + 1) inds + go [] i inds = inds