diff --git a/ChangeLog.md b/ChangeLog.md new file mode 100644 --- /dev/null +++ b/ChangeLog.md @@ -0,0 +1,11 @@ + +# Markov chain Monte Carlo sampling - ChangeLog + + +## 0.1.3 + +Many changes; notably it is now possible to continue a Markov chain run. + + +## Unreleased changes + diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,30 @@ + +# Markov chain Monte Carlo + +

+ +Sample from a posterior using Markov chain Monte Carlo. + +At the moment, the library is tailored to the Metropolis-Hastings algorithm +since it covers most use cases. However, implementation of more algorithms is +planned in the future. + + +## Documentation + +The source code contains detailed documentation about general concepts as well +as specific functions. + + +## Examples + +Have a look at the [example MCMC analyses](https://github.com/dschrempf/mcmc/tree/master/mcmc-examples). They can be built with [Stack](https://docs.haskellstack.org/en/stable/README/) and are +attached to this repository. + + git clone https://github.com/dschrempf/mcmc.git + stack build + +For example, estimate the [accuracy of an archer](https://github.com/dschrempf/mcmc/blob/master/mcmc-examples/Archery/Archery.hs) with + + stack exec archery + diff --git a/bench/Bench.hs b/bench/Bench.hs new file mode 100644 --- /dev/null +++ b/bench/Bench.hs @@ -0,0 +1,27 @@ +-- | +-- Module : Main +-- Copyright : (c) Dominik Schrempf 2020 +-- License : GPL-3.0-or-later +-- +-- Maintainer : dominik.schrempf@gmail.com +-- Stability : unstable +-- Portability : portable +-- +-- Creation date: Fri May 8 17:27:03 2020. +module Main + ( main, + ) +where + +import Criterion.Main +import Normal +import Poisson +import System.Random.MWC + +main :: IO () +main = do + g <- create + defaultMain + [ bench "Normal" $ nfIO (normalBench g), + bench "Poisson" $ nfIO (poissonBench g) + ] diff --git a/bench/Normal.hs b/bench/Normal.hs new file mode 100644 --- /dev/null +++ b/bench/Normal.hs @@ -0,0 +1,63 @@ +-- | +-- Module : Normal +-- Description : Benchmark Metropolis-Hastings algorithm +-- Copyright : (c) Dominik Schrempf 2020 +-- License : GPL-3.0-or-later +-- +-- Maintainer : dominik.schrempf@gmail.com +-- Stability : unstable +-- Portability : portable +-- +-- Creation date: Wed May 6 00:10:11 2020. +module Normal + ( normalBench, + ) +where + +import Control.Monad +import Mcmc +import Numeric.Log as L +import Statistics.Distribution hiding + ( mean, + stdDev, + ) +import Statistics.Distribution.Normal +import System.Random.MWC + +trueMean :: Double +trueMean = 5 + +trueStdDev :: Double +trueStdDev = 4 + +lh :: Double -> Log Double +lh = Exp . logDensity (normalDistr trueMean trueStdDev) + +moveCycle :: Cycle Double +moveCycle = + fromList + [ slideSymmetric "small" 5 id 0.1 True, + slideSymmetric "medium" 2 id 1.0 True, + slideSymmetric "large" 2 id 5.0 True, + slide "skewed" 1 id 1.0 4.0 True + ] + +monStd :: MonitorStdOut Double +monStd = monitorStdOut [monitorRealFloat "mu" id] 200 + +mon :: Monitor Double +mon = Monitor monStd [] [] + +nBurn :: Maybe Int +nBurn = Just 2000 + +nAutoTune :: Maybe Int +nAutoTune = Just 200 + +nIter :: Int +nIter = 20000 + +normalBench :: GenIO -> IO () +normalBench g = do + let s = noSave $ status "Normal" (const 1) lh moveCycle mon 0 nBurn nAutoTune nIter g + void $ mh s diff --git a/bench/Poisson.hs b/bench/Poisson.hs new file mode 100644 --- /dev/null +++ b/bench/Poisson.hs @@ -0,0 +1,86 @@ +-- | +-- Module : Poisson +-- Description : Poisson regression model for airline fatalities +-- Copyright : (c) Dominik Schrempf 2020 +-- License : GPL-3.0-or-later +-- +-- Maintainer : dominik.schrempf@gmail.com +-- Stability : unstable +-- Portability : portable +-- +-- Creation date: Fri May 8 13:30:51 2020. +-- +-- See https://revbayes.github.io/tutorials/mcmc/poisson.html. +module Poisson + ( poissonBench, + ) +where + +import Control.Monad +import Lens.Micro +import Mcmc +import Numeric.Log hiding (sum) +import Statistics.Distribution hiding + ( mean, + stdDev, + ) +import Statistics.Distribution.Poisson +import System.Random.MWC + +type I = (Double, Double) + +fatalities :: [Int] +fatalities = [24, 25, 31, 31, 22, 21, 26, 20, 16, 22] + +normalizedYears :: [Double] +normalizedYears = map (subtract m) ys + where + ys = [1976.0 .. 1985.0] + m = sum ys / fromIntegral (length ys) + +f :: Int -> Double -> I -> Log Double +f ft yr (alpha, beta) = Exp $ logProbability (poisson l) (fromIntegral ft) + where + l = exp $ alpha + beta * yr + +lh :: I -> Log Double +lh x = + product [f ft yr x | (ft, yr) <- zip fatalities normalizedYears] + +moveAlpha :: Move I +moveAlpha = slideSymmetric "alpha" 2 _1 0.2 False + +moveBeta :: Move I +moveBeta = slideSymmetric "beta" 1 _2 0.2 False + +moveCycle :: Cycle I +moveCycle = fromList [moveAlpha, moveBeta] + +initial :: I +initial = (0, 0) + +monAlpha :: MonitorParameter I +monAlpha = monitorRealFloat "alpha" _1 + +monBeta :: MonitorParameter I +monBeta = monitorRealFloat "beta" _2 + +monStd :: MonitorStdOut I +monStd = monitorStdOut [monAlpha, monBeta] 150 + +mon :: Monitor I +mon = Monitor monStd [] [] + +nBurn :: Maybe Int +nBurn = Just 2000 + +nAutoTune :: Maybe Int +nAutoTune = Just 200 + +nIter :: Int +nIter = 10000 + +poissonBench :: GenIO -> IO () +poissonBench g = do + let s = noSave $ status "Poisson" (const 1) lh moveCycle mon initial nBurn nAutoTune nIter g + void $ mh s diff --git a/mcmc.cabal b/mcmc.cabal new file mode 100644 --- /dev/null +++ b/mcmc.cabal @@ -0,0 +1,112 @@ +cabal-version: 2.2 + +-- This file has been generated from package.yaml by hpack version 0.34.2. +-- +-- see: https://github.com/sol/hpack + +name: mcmc +version: 0.1.3 +synopsis: Sample from a posterior using Markov chain Monte Carlo +description: Please see the README on GitHub at +category: Math, Statistics +homepage: https://github.com/dschrempf/mcmc#readme +bug-reports: https://github.com/dschrempf/mcmc/issues +author: Dominik Schrempf +maintainer: dominik.schrempf@gmail.com +copyright: Dominik Schrempf (2020) +license: GPL-3.0-or-later +build-type: Simple +extra-source-files: + README.md + ChangeLog.md + +source-repository head + type: git + location: https://github.com/dschrempf/mcmc + +library + exposed-modules: + Mcmc + Mcmc.Item + Mcmc.Mcmc + Mcmc.Metropolis + Mcmc.Monitor + Mcmc.Monitor.Log + Mcmc.Monitor.Parameter + Mcmc.Monitor.ParameterBatch + Mcmc.Monitor.Time + Mcmc.Move + Mcmc.Move.Generic + Mcmc.Move.Scale + Mcmc.Move.Slide + Mcmc.Save + Mcmc.Status + Mcmc.Tools.Shuffle + Mcmc.Trace + other-modules: + Paths_mcmc + autogen-modules: + Paths_mcmc + hs-source-dirs: + src + ghc-options: -Wall + build-depends: + aeson + , base >=4.7 && <5 + , bytestring + , containers + , directory + , log-domain + , microlens + , mwc-random + , statistics + , text + , time + , transformers + , vector + , zlib + default-language: Haskell2010 + +test-suite mcmc-test + type: exitcode-stdio-1.0 + main-is: Spec.hs + other-modules: + Mcmc.Move.SlideSpec + Mcmc.SaveSpec + Paths_mcmc + hs-source-dirs: + test + ghc-options: -Wall -threaded -rtsopts -with-rtsopts=-N + build-depends: + QuickCheck + , base >=4.7 && <5 + , hspec + , hspec-discover + , log-domain + , mcmc + , mwc-random + , statistics + , vector + default-language: Haskell2010 + +benchmark mcmc-bench + type: exitcode-stdio-1.0 + main-is: Bench.hs + other-modules: + Normal + Poisson + Paths_mcmc + hs-source-dirs: + bench + ghc-options: -Wall -threaded -rtsopts -with-rtsopts=-N + build-depends: + base >=4.7 && <5 + , criterion + , log-domain + , mcmc + , microlens + , mwc-random + , statistics + , text + , vector + default-language: Haskell2010 diff --git a/src/Mcmc.hs b/src/Mcmc.hs new file mode 100644 --- /dev/null +++ b/src/Mcmc.hs @@ -0,0 +1,145 @@ +{-# LANGUAGE RankNTypes #-} + +-- | +-- Module : Mcmc +-- Description : Markov chain Monte Carlo algorithms, batteries included +-- Copyright : (c) Dominik Schrempf 2020 +-- License : GPL-3.0-or-later +-- +-- Maintainer : dominik.schrempf@gmail.com +-- Stability : unstable +-- Portability : portable +-- +-- Creation date: Tue May 5 18:01:15 2020. +-- +-- A short introduction to update mechanisms using the Metropolis-Hastings +-- algorithm (see Geyer, C. J., 2011; Introduction to Markov Chain Monte Carlo. In +-- Handbook of Markov Chain Monte Carlo (pp. 45), Chapman \& Hall/CRC). +-- +-- For examples, please see +-- [mcmc-examples](https://github.com/dschrempf/mcmc/tree/master/mcmc-examples). +-- +-- The import of this module alone should cover most use cases. +module Mcmc + ( -- * Moves + + -- | A 'Move' is an instruction about how to advance a given Markov chain so + -- that it possibly reaches a new state. That is, 'Move's specify how the + -- chain traverses the state space. As far as this MCMC library is + -- concerned, 'Move's are /elementary updates/ in that they cannot be + -- decomposed into smaller updates. + -- + -- 'Move's can be combined to form composite updates, a technique often + -- referred to as /composition/. On the other hand, /mixing/ (used in the + -- sense of mixture models) is the random choice of a 'Move' (or a + -- composition of 'Move's) from a given set. + -- + -- The __composition__ and __mixture__ of 'Move's allows specification of + -- nearly all MCMC algorithms involving a single chain (i.e., population + -- methods such as particle filters are excluded). In particular, Gibbs + -- samplers of all sorts can be specified using this procedure. + -- + -- This library enables composition and mixture of 'Move's via the 'Cycle' + -- data type. Essentially, a 'Cycle' is a set of 'Move's. The chain advances + -- after the completion of each 'Cycle', which is called an __iteration__, + -- and the iteration counter is increased by one. + -- + -- The 'Move's in a 'Cycle' can be executed in the given order or in a + -- random sequence which allows, for example, specification of a fixed scan + -- Gibbs sampler, or a random sequence scan Gibbs sampler, respectively. + -- + -- Note that it is of utter importance that the given 'Cycle' enables + -- traversal of the complete state space. Otherwise, the Markov chain will + -- not converge to the correct stationary posterior distribution. + -- + -- Moves are named according to what they do, i.e., how they change the + -- state of a Markov chain, and not according to the intrinsically used + -- probability distributions. For example, 'slideSymmetric' is a sliding + -- move. Under the hood, it uses the normal distribution with mean zero and + -- given variance. The sampled variate is added to the current value of the + -- variable (hence, the name slide). The same nomenclature is used by + -- RevBayes [1]. The probability distributions and intrinsic properties of a + -- specific move are specified in detail in the documentation. + -- + -- The other method, which is used intrinsically, is more systematic, but + -- also a little bit more complicated: we separate between the proposal + -- distribution and how the state is affected. And here, I am not only + -- referring to the accessor (i.e., the lens), but also to the operator + -- (addition, multiplication, any other binary operator). For example, the + -- sliding move (without tuning information) is implemented as + -- + -- @ + -- slideSimple :: Lens' a Double -> Double -> Double -> Double -> MoveSimple a + -- slideSimple l m s t = moveGenericContinuous l (normalDistr m (s * t)) (+) (-) + -- @ + -- + -- This specification is more involved. Especially since we need to know the + -- probability of jumping back, and so we need to know the inverse operator. + -- However, it also allows specification of new moves with great ease. + -- + -- [1] Höhna, S., Landis, M. J., Heath, T. A., Boussau, B., Lartillot, N., Moore, + -- B. R., Huelsenbeck, J. P., …, Revbayes: bayesian phylogenetic inference using + -- graphical models and an interactive model-specification language, Systematic + -- Biology, 65(4), 726–736 (2016). http://dx.doi.org/10.1093/sysbio/syw021 + module Mcmc.Move, + module Mcmc.Move.Slide, + module Mcmc.Move.Scale, + + -- * Initialization + + -- | The 'Status' contains all information to run an MCMC chain. It is + -- constructed using the function 'status'. + -- + -- The 'Status' of a Markov chain includes information about current state + -- ('Item') and iteration, the history of the chain ('Trace'), the 'Acceptance' + -- ratios, and the random number generator. + -- + -- Further, the 'Status' includes auxiliary variables and functions such as + -- the prior and likelihood functions, instructions to move around the state + -- space (see above) and to monitor the MCMC run, as well as some auxiliary + -- information. + module Mcmc.Status, + module Mcmc.Item, + module Mcmc.Trace, + + -- * Monitor + + -- | A 'Monitor' describes which part of the Markov chain should be logged + -- and where. There are three different types: + -- - 'MonitorStdOut': Log to standard output. + -- - 'MonitorFile': Log to a file. + -- - 'MonitorBatch': Log summary statistics such as the mean of the last + -- - states to a file. + module Mcmc.Monitor, + module Mcmc.Monitor.Parameter, + module Mcmc.Monitor.ParameterBatch, + + -- * Algorithms + + -- | At the moment, the library is tailored to the Metropolis-Hastings + -- algorithm ('mh') since it covers most use cases. However, implementation + -- of more algorithms is planned in the future. + module Mcmc.Metropolis, + + -- * Misc + pzero, + module Mcmc.Save, + ) +where + +import Mcmc.Item +import Mcmc.Metropolis +import Mcmc.Monitor +import Mcmc.Monitor.Parameter +import Mcmc.Monitor.ParameterBatch +import Mcmc.Move +import Mcmc.Move.Scale +import Mcmc.Move.Slide +import Mcmc.Save +import Mcmc.Status +import Mcmc.Trace +import Numeric.Log + +-- | Because we need a probability of zero for likelihoods of really bad moves. +pzero :: Fractional a => Log a +pzero = Exp $ - (1 / 0) diff --git a/src/Mcmc/Item.hs b/src/Mcmc/Item.hs new file mode 100644 --- /dev/null +++ b/src/Mcmc/Item.hs @@ -0,0 +1,54 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} + +-- | +-- Module : Mcmc.Item +-- Description : Links of Markov chains +-- Copyright : (c) Dominik Schrempf 2020 +-- License : GPL-3.0-or-later +-- +-- Maintainer : dominik.schrempf@gmail.com +-- Stability : unstable +-- Portability : portable +-- +-- Creation date: Wed May 20 09:10:27 2020. +module Mcmc.Item + ( Item (..), + ) +where + +import Data.Aeson +import Numeric.Log + +instance ToJSON a => ToJSON (Log a) where + toJSON (Exp x) = toJSON x + toEncoding (Exp x) = toEncoding x + +instance FromJSON a => FromJSON (Log a) where + parseJSON v = Exp <$> parseJSON v + +-- | An 'Item', or link of the Markov chain. For reasons of computational +-- efficiency, each state is associated with the corresponding prior and +-- likelihood. +data Item a + = Item + { -- | The current state in the state space @a@. + state :: a, + -- | The current prior. + prior :: Log Double, + -- | The current likelihood. + likelihood :: Log Double + } + deriving (Eq, Ord, Show, Read) + +instance ToJSON a => ToJSON (Item a) where + toJSON (Item x p l) = object ["s" .= x, "p" .= p, "l" .= l] + toEncoding (Item x p l) = pairs ("s" .= x <> "p" .= p <> "l" .= l) + +instance FromJSON a => FromJSON (Item a) where + parseJSON = withObject "Item" $ + \v -> + Item + <$> v .: "s" + <*> v .: "p" + <*> v .: "l" diff --git a/src/Mcmc/Mcmc.hs b/src/Mcmc/Mcmc.hs new file mode 100644 --- /dev/null +++ b/src/Mcmc/Mcmc.hs @@ -0,0 +1,147 @@ +{-# LANGUAGE OverloadedStrings #-} + +-- | +-- Module : Mcmc.Mcmc +-- Description : Mcmc helpers +-- Copyright : (c) Dominik Schrempf, 2020 +-- License : GPL-3.0-or-later +-- +-- Maintainer : dominik.schrempf@gmail.com +-- Stability : unstable +-- Portability : portable +-- +-- Creation date: Fri May 29 10:19:45 2020. +module Mcmc.Mcmc + ( Mcmc, + mcmcAutotune, + mcmcSummarizeCycle, + mcmcInit, + mcmcReport, + mcmcMonitorHeader, + mcmcMonitorExec, + mcmcClose, + ) +where + +import Control.Monad +import Control.Monad.IO.Class +import Control.Monad.Trans.State.Strict +import Data.Aeson +import Data.Maybe +import qualified Data.Text.Lazy.IO as T +import Data.Time.Clock +import Data.Time.Format +import Mcmc.Monitor +import Mcmc.Monitor.Time +import Mcmc.Move +import Mcmc.Save +import Mcmc.Status +import Prelude hiding (cycle) + +-- | An Mcmc state transformer; usually fiddling around with this type is not +-- required, but it is used by the different inference algorithms. +type Mcmc a = StateT (Status a) IO + +-- | Auto tune the 'Move's in the 'Cycle' of the chain. See 'autotune'. +mcmcAutotune :: Int -> Mcmc a () +mcmcAutotune t = do + s <- get + let a = acceptance s + c = cycle s + c' = autotuneC t a c + put $ s {cycle = c'} + +-- | Print short summary of 'Move's in 'Cycle'. See 'summarizeCycle'. +mcmcSummarizeCycle :: Maybe Int -> Mcmc a () +mcmcSummarizeCycle Nothing = do + c <- gets cycle + liftIO $ T.putStr $ summarizeCycle Nothing c +mcmcSummarizeCycle (Just n) = do + a <- gets acceptance + c <- gets cycle + liftIO $ T.putStr $ summarizeCycle (Just (n, a)) c + +fTime :: FormatTime t => t -> String +fTime = formatTime defaultTimeLocale "%B %-e, %Y, at %H:%M %P, %Z." + +-- | Set the total number of iterations, the current time and open the +-- 'Monitor's of the chain. See 'mOpen'. +mcmcInit :: Mcmc a () +mcmcInit = do + t <- liftIO getCurrentTime + liftIO $ putStrLn $ "-- Start time: " <> fTime t + s <- get + let m = monitor s + n = iteration s + nm = name s + m' <- if n == 0 + then liftIO $ mOpen nm m + else liftIO $ mAppend nm m + put $ s {monitor = m', start = Just (n, t)} + +-- | Report what is going to be done. +mcmcReport :: ToJSON a => Mcmc a () +mcmcReport = do + s <- get + let b = burnInIterations s + t = autoTuningPeriod s + n = iterations s + case b of + Just b' -> liftIO $ putStrLn $ "-- Burn in for " <> show b' <> " iterations." + Nothing -> return () + case t of + Just t' -> + liftIO $ putStrLn $ + "-- Auto tune every " + <> show t' + <> " iterations (during burn in only)." + Nothing -> return () + liftIO $ putStrLn $ "-- Run chain for " <> show n <> " iterations." + liftIO $ putStrLn "-- Initial state." + mcmcMonitorHeader + mcmcMonitorExec + +-- | Print header line of 'Monitor' (only standard output). +mcmcMonitorHeader :: Mcmc a () +mcmcMonitorHeader = do + m <- gets monitor + liftIO $ mHeader m + +-- Save the status of an MCMC run. See 'saveStatus'. +mcmcSave :: ToJSON a => Mcmc a () +mcmcSave = do + s <- get + liftIO $ if save s + then do putStrLn "-- Save Markov chain. For long chains, this may take a while." + saveStatus (name s <> ".mcmc") s + putStrLn "-- Done saving Markov chain." + else putStrLn "-- Do not save the Markov chain." + +-- | Execute the 'Monitor's of the chain. See 'mExec'. +mcmcMonitorExec :: ToJSON a => Mcmc a () +mcmcMonitorExec = do + s <- get + let i = iteration s + j = iterations s + fromMaybe 0 (burnInIterations s) + m = monitor s + st = fromMaybe (error "mcmcMonitorExec: Starting state and time not set.") (start s) + tr = trace s + liftIO $ mExec i st tr j m + +-- | Close the 'Monitor's of the chain. See 'mClose'. +mcmcClose :: ToJSON a => Mcmc a () +mcmcClose = do + s <- get + let n = iterations s + mcmcSummarizeCycle (Just n) + liftIO $ putStrLn "-- Metropolis-Hastings sampler finished." + let m = monitor s + m' <- liftIO $ mClose m + put $ s {monitor = m'} + mcmcSave + t <- liftIO getCurrentTime + let rt = case start s of + Nothing -> error "mcmcClose: Start time not set." + Just (_, st) -> t `diffUTCTime` st + liftIO $ T.putStrLn $ "-- Wall clock run time: " <> renderDuration rt <> "." + liftIO $ putStrLn $ "-- End time: " <> fTime t diff --git a/src/Mcmc/Metropolis.hs b/src/Mcmc/Metropolis.hs new file mode 100644 --- /dev/null +++ b/src/Mcmc/Metropolis.hs @@ -0,0 +1,180 @@ +{-# LANGUAGE BangPatterns #-} + +-- | +-- Module : Mcmc.Metropolis +-- Description : Metropolis-Hastings at its best +-- Copyright : (c) Dominik Schrempf 2020 +-- License : GPL-3.0-or-later +-- +-- Maintainer : dominik.schrempf@gmail.com +-- Stability : unstable +-- Portability : portable +-- +-- Creation date: Tue May 5 20:11:30 2020. +module Mcmc.Metropolis + ( mh, + mhContinue, + ) +where + +import Control.Monad +import Control.Monad.IO.Class +import Control.Monad.Trans.State.Strict +import Data.Aeson +import Data.Maybe +import Mcmc.Item +import Mcmc.Mcmc +import Mcmc.Move +import Mcmc.Status +import Mcmc.Tools.Shuffle +import Mcmc.Trace +import Numeric.Log +import System.Random.MWC +import Prelude hiding (cycle) + +-- For non-symmetric moves. +mhRatio :: Log Double -> Log Double -> Log Double -> Log Double -> Log Double +mhRatio lX lY qXY qYX = lY * qYX / lX / qXY +{-# INLINE mhRatio #-} + +-- For symmetric moves. +mhRatioSymmetric :: Log Double -> Log Double -> Log Double +mhRatioSymmetric lX lY = lY / lX +{-# INLINE mhRatioSymmetric #-} + +mhMove :: Move a -> Mcmc a () +mhMove m = do + let p = mvSample $ mvSimple m + mq = mvDensity $ mvSimple m + s <- get + let (Item x pX lX) = item s + pF = priorF s + lF = likelihoodF s + a = acceptance s + g = generator s + -- 1. Sample new state. + !y <- liftIO $ p x g + -- 2. Calculate Metropolis-Hastings ratio. + let !pY = pF y + !lY = lF y + !r = case mq of + Nothing -> mhRatioSymmetric (pX * lX) (pY * lY) + Just q -> mhRatio (pX * lX) (pY * lY) (q x y) (q y x) + -- 3. Accept or reject. + if ln r >= 0.0 + then put $ s {item = Item y pY lY, acceptance = pushA m True a} + else do + b <- uniform g + if b < exp (ln r) + then put $ s {item = Item y pY lY, acceptance = pushA m True a} + else put $ s {acceptance = pushA m False a} + +-- TODO: Split the generator here. See SaveSpec -> mhContinue. + +-- Replicate 'Move's according to their weights and shuffle them. +getNCycles :: Cycle a -> Int -> GenIO -> IO [[Move a]] +getNCycles c = shuffleN mvs + where + !mvs = concat [replicate (mvWeight m) m | m <- fromCycle c] + +-- Run one iterations; perform all moves in a Cycle. +mhIter :: ToJSON a => [Move a] -> Mcmc a () +mhIter mvs = do + mapM_ mhMove mvs + s <- get + let i = item s + t = trace s + n = iteration s + put $ s {trace = pushT i t, iteration = succ n} + mcmcMonitorExec + +-- Run N iterations. +mhNIter :: ToJSON a => Int -> Mcmc a () +mhNIter n = do + c <- gets cycle + g <- gets generator + cycles <- liftIO $ getNCycles c n g + forM_ cycles mhIter + +-- Burn in and auto tune. +mhBurnInN :: ToJSON a => Int -> Maybe Int -> Mcmc a () +mhBurnInN b (Just t) + | t <= 0 = error "mhBurnInN: Auto tuning period smaller equal 0." + | b > t = mhNIter t >> mcmcAutotune t >> mhBurnInN (b - t) (Just t) + | otherwise = mhNIter b >> mcmcAutotune b +mhBurnInN b Nothing = mhNIter b + +-- Initialize burn in for given number of iterations. +mhBurnIn :: ToJSON a => Int -> Maybe Int -> Mcmc a () +mhBurnIn b t + | b < 0 = error "mhBurnIn: Negative number of burn in iterations." + | b == 0 = return () + | otherwise = do + liftIO $ putStrLn $ "-- Burn in for " <> show b <> " cycles." + mcmcMonitorHeader + mhBurnInN b t + liftIO $ putStrLn "-- Burn in finished." + case t of + Nothing -> return () + Just _ -> mcmcSummarizeCycle t + s <- get + let a = acceptance s + put $ s {acceptance = resetA a} + +-- Run for given number of iterations. +mhRun :: ToJSON a => Int -> Mcmc a () +mhRun n = do + liftIO $ putStrLn $ "-- Run chain for " <> show n <> " iterations." + mcmcMonitorHeader + mhNIter n + +mhT :: ToJSON a => Mcmc a () +mhT = do + s <- get + liftIO $ putStrLn "-- Start of Metropolis-Hastings sampler." + mcmcInit + mcmcReport + mcmcSummarizeCycle Nothing + let b = fromMaybe 0 (burnInIterations s) + mhBurnIn b (autoTuningPeriod s) + let n = iterations s + mhRun n + mcmcClose + +mhContinueT :: ToJSON a => Int -> Mcmc a () +mhContinueT dn = do + liftIO $ putStrLn "-- Continue Metropolis-Hastings sampler." + liftIO $ putStrLn $ "-- Run chain for " <> show dn <> " additional iterations." + s <- get + let n = iterations s + put s {iterations = n + dn} + mcmcInit + mcmcSummarizeCycle Nothing + mhRun dn + mcmcClose + +-- | Continue a Markov chain for a given number of Metropolis-Hastings steps. +mhContinue :: + ToJSON a => + -- | Additional number of Metropolis-Hastings steps. + Int -> + -- | Loaded status of the Markov chain. + Status a -> + IO (Status a) +mhContinue dn + | dn <= 0 = error "mhContinue: The number of iterations is zero or negative." + | otherwise = execStateT $ mhContinueT dn + +-- | Run a Markov chain for a given number of Metropolis-Hastings steps. +mh :: + ToJSON a => + -- | Initial (or last) status of the Markov chain. + Status a -> + IO (Status a) +mh s = do + let m = iteration s + if m == 0 + then execStateT mhT s + else do + putStrLn "To continue a Markov chain run, please use 'mhContinue'." + error $ "mh: Current iteration " ++ show m ++ " is non-zero." diff --git a/src/Mcmc/Monitor.hs b/src/Mcmc/Monitor.hs new file mode 100644 --- /dev/null +++ b/src/Mcmc/Monitor.hs @@ -0,0 +1,355 @@ +{-# LANGUAGE OverloadedStrings #-} + +-- | +-- Module : Mcmc.Monitor +-- Description : Monitor a Markov chain +-- Copyright : (c) Dominik Schrempf, 2020 +-- License : GPL-3.0-or-later +-- +-- Maintainer : dominik.schrempf@gmail.com +-- Stability : unstable +-- Portability : portable +-- +-- Creation date: Thu May 21 14:35:11 2020. +module Mcmc.Monitor + ( -- * Create monitors + Monitor (..), + MonitorStdOut, + monitorStdOut, + MonitorFile, + monitorFile, + MonitorBatch, + monitorBatch, + + -- * Use monitor + mOpen, + mAppend, + mHeader, + mExec, + mClose, + ) +where + +import Control.Monad +import Data.Int +import qualified Data.Text.Lazy as T +import Data.Text.Lazy (Text) +import qualified Data.Text.Lazy.Builder as T +import qualified Data.Text.Lazy.IO as T +import Data.Time.Clock +import Mcmc.Item +import Mcmc.Monitor.Log +import Mcmc.Monitor.Parameter +import Mcmc.Monitor.ParameterBatch +import Mcmc.Monitor.Time +import Mcmc.Trace +import Numeric.Log +import System.Directory +import System.IO +import Prelude hiding (sum) + +-- | A 'Monitor' describes which part of the Markov chain should be logged and +-- where. Further, they allow output of summary statistics per iteration in a +-- flexible way. +data Monitor a + = Monitor + { -- | Monitor writing to standard output. + mStdOut :: MonitorStdOut a, + -- | Monitors writing to files. + mFiles :: [MonitorFile a], + -- | Monitors calculating batch means and + -- writing to files. + mBatches :: [MonitorBatch a] + } + +-- | Monitor to standard output. +data MonitorStdOut a + = MonitorStdOut + { msParams :: [MonitorParameter a], + msPeriod :: Int + } + +-- | Monitor to standard output. +monitorStdOut :: + -- | Instructions about which parameters to log. + [MonitorParameter a] -> + -- | Logging period. + Int -> + MonitorStdOut a +monitorStdOut ps p + | p < 1 = error "monitorStdOut: Monitor period has to be 1 or larger." + | otherwise = MonitorStdOut ps p + +msIWidth :: Int64 +msIWidth = 12 + +msWidth :: Int64 +msWidth = 22 + +msRenderRow :: [Text] -> Text +msRenderRow xs = T.justifyRight msIWidth ' ' (head xs) <> T.concat vals + where + vals = map (T.justifyRight msWidth ' ') (tail xs) + +msHeader :: MonitorStdOut a -> IO () +msHeader m = T.hPutStr stdout $ T.unlines [row, sep] + where + row = + msRenderRow $ + ["Iteration", "Log-Prior", "Log-Likelihood", "Log-Posterior"] + ++ nms + ++ ["Runtime", "ETA"] + sep = " " <> T.replicate (T.length row - 3) "─" + nms = [T.pack $ mpName p | p <- msParams m] + +msExec :: + Int -> + Item a -> + (Int, UTCTime) -> + Int -> + MonitorStdOut a -> + IO () +msExec i (Item x p l) (ss, st) j m + | i `mod` msPeriod m /= 0 = + return () + | otherwise = do + ct <- getCurrentTime + let dt = ct `diffUTCTime` st + timePerIter = dt / fromIntegral (i - ss) + eta = if i < 10 + then "" + else renderDuration $ timePerIter * fromIntegral (j - i) + T.hPutStrLn stdout + $ msRenderRow + $ [T.pack (show i), renderLog p, renderLog l, renderLog (p * l)] + ++ [T.toLazyText $ mpFunc mp x | mp <- msParams m] + ++ [renderDuration dt , eta] + +-- | Monitor to a file. +data MonitorFile a + = MonitorFile + { mfName :: String, + mfHandle :: Maybe Handle, + mfParams :: [MonitorParameter a], + mfPeriod :: Int + } + +-- XXX: The file monitor also includes iteration, prior, likelihood, and +-- posterior. What if I want to log trees; or other complex objects? In this +-- case, we need a simpler monitor to a file. + +-- | Monitor parameters to a file. +monitorFile :: + -- | Name; used as part of the file name. + String -> + -- | Instructions about which parameters to log. + [MonitorParameter a] -> + -- | Logging period. + Int -> + MonitorFile a +monitorFile n ps p + | p < 1 = error "monitorFile: Monitor period has to be 1 or larger." + | otherwise = MonitorFile n Nothing ps p + +mfRenderRow :: [Text] -> Text +mfRenderRow = T.intercalate "\t" + +mfOpen :: String -> MonitorFile a -> IO (MonitorFile a) +mfOpen n m = do + h <- openFile (n <> mfName m <> ".monitor") WriteMode + hSetBuffering h LineBuffering + return $ m {mfHandle = Just h} + +mfAppend :: String -> MonitorFile a -> IO (MonitorFile a) +mfAppend n m = do + let fn = n <> mfName m <> ".monitor" + fe <- doesFileExist fn + if fe + then do h <- openFile fn AppendMode + hSetBuffering h LineBuffering + return $ m {mfHandle = Just h} + else error $ "mfAppend: Monitor file does not exist: " ++ fn ++ "." + +mfHeader :: MonitorFile a -> IO () +mfHeader m = case mfHandle m of + Nothing -> + error $ + "mfHeader: No handle available for monitor with name " + <> mfName m + <> "." + Just h -> + T.hPutStrLn h + $ mfRenderRow + $ ["Iteration", "Log-Prior", "Log-Likelihood", "Log-Posterior"] + ++ [T.pack $ mpName p | p <- mfParams m] + +mfExec :: + Int -> + Item a -> + MonitorFile a -> + IO () +mfExec i (Item x p l) m + | i `mod` mfPeriod m /= 0 = return () + | otherwise = case mfHandle m of + Nothing -> + error $ + "mfExec: No handle available for monitor with name " + <> mfName m + <> "." + Just h -> + T.hPutStrLn h + $ mfRenderRow + $ T.pack (show i) + : renderLog p + : renderLog l + : renderLog (p * l) + : [T.toLazyText $ mpFunc mp x | mp <- mfParams m] + +mfClose :: MonitorFile a -> IO () +mfClose m = case mfHandle m of + Just h -> hClose h + Nothing -> error $ "mfClose: File was not opened for monitor " <> mfName m <> "." + +-- | Monitor to a file, but calculate batch means for the given batch size. +-- +-- XXX: Batch monitors are slow at the moment because the monitored parameter +-- has to be extracted from the state for each iteration. +data MonitorBatch a + = MonitorBatch + { mbName :: String, + mbHandle :: Maybe Handle, + mbParams :: [MonitorParameterBatch a], + mbSize :: Int + } + +-- XXX: The batch monitor also includes iteration, prior, likelihood, and +-- posterior. What if I want to log trees; or other complex objects? In this +-- case, we need a simpler monitor to a file. + +-- | Monitor parameters to a file, see 'MonitorBatch'. +monitorBatch :: + -- | Name; used as part of the file name. + String -> + -- | Instructions about which parameters to log + -- and how to calculate the batch means. + [MonitorParameterBatch a] -> + -- | Batch size. + Int -> + MonitorBatch a +monitorBatch n ps p + | p < 2 = error "monitorBatch: Batch size has to be 2 or larger." + | otherwise = MonitorBatch n Nothing ps p + +mbOpen :: String -> MonitorBatch a -> IO (MonitorBatch a) +mbOpen n m = do + h <- openFile (n <> mbName m <> ".batch") WriteMode + hSetBuffering h LineBuffering + return $ m {mbHandle = Just h} + +mbAppend :: String -> MonitorBatch a -> IO (MonitorBatch a) +mbAppend n m = do + let fn = n <> mbName m <> ".batch" + fe <- doesFileExist fn + if fe + then do h <- openFile fn AppendMode + hSetBuffering h LineBuffering + return $ m {mbHandle = Just h} + else error $ "mbAppend: Monitor file does not exist: " ++ fn ++ "." + +mbHeader :: MonitorBatch a -> IO () +mbHeader m = case mbHandle m of + Nothing -> + error $ + "mbHeader: No handle available for batch monitor with name " + <> mbName m + <> "." + Just h -> + T.hPutStrLn h + $ mfRenderRow + $ ["Iteration", "Mean log-Prior", "Mean log-Likelihood", "Mean log-Posterior"] + ++ [T.pack $ mbpName mbp | mbp <- mbParams m] + +mean :: [Log Double] -> Log Double +mean xs = sum xs / fromIntegral (length xs) + +mbExec :: + Int -> + Trace a -> + MonitorBatch a -> + IO () +mbExec i t' m + | (i `mod` mbSize m /= 0) || (i == 0) = return () + | otherwise = case mbHandle m of + Nothing -> + error $ + "mbExec: No handle available for batch monitor with name " + <> mbName m + <> "." + Just h -> + T.hPutStrLn h + $ mfRenderRow + $ T.pack (show i) + : renderLog mlps + : renderLog mlls + : renderLog mlos + : [T.toLazyText $ mbpFunc mbp (map state t) | mbp <- mbParams m] + where + t = takeT (mbSize m) t' + lps = map prior t + lls = map likelihood t + los = zipWith (*) lps lls + mlps = mean lps + mlls = mean lls + mlos = mean los + +mbClose :: MonitorBatch a -> IO () +mbClose m = case mbHandle m of + Just h -> hClose h + Nothing -> error $ "mfClose: File was not opened for batch monitor: " <> mbName m <> "." + +-- | Open the files associated with the 'Monitor'. +mOpen :: String -> Monitor a -> IO (Monitor a) +mOpen n (Monitor s fs bs) = do + fs' <- mapM (mfOpen n) fs + mapM_ mfHeader fs' + bs' <- mapM (mbOpen n) bs + mapM_ mbHeader bs' + return $ Monitor s fs' bs' + +-- | Open the files associated with the 'Monitor' in append mode. +mAppend :: String -> Monitor a -> IO (Monitor a) +mAppend n (Monitor s fs bs) = do + fs' <- mapM (mfAppend n) fs + bs' <- mapM (mbAppend n) bs + return $ Monitor s fs' bs' + +-- | Print header line of 'Monitor' (standard output only). +mHeader :: Monitor a -> IO () +mHeader (Monitor s _ _) = msHeader s + +-- | Execute monitors; print status information to standard output and files. +mExec :: + -- | Iteration. + Int -> + -- | Starting state and time. + (Int, UTCTime) -> + -- | Trace of Markov chain. + Trace a -> + -- | Total number of iterations; to calculate ETA. + Int -> + -- | The monitor. + Monitor a -> + IO () +mExec i t xs j (Monitor s fs bs) = do + msExec i (headT xs) t j s + mapM_ (mfExec i $ headT xs) fs + mapM_ (mbExec i xs) bs + +-- | Close the files associated with the 'Monitor'. +mClose :: Monitor a -> IO (Monitor a) +mClose m@(Monitor _ fms bms) = do + mapM_ mfClose fms + mapM_ mbClose bms + let fms' = map (\fm -> fm {mfHandle = Nothing}) fms + let bms' = map (\bm -> bm {mbHandle = Nothing}) bms + return m {mFiles = fms', mBatches = bms'} diff --git a/src/Mcmc/Monitor/Log.hs b/src/Mcmc/Monitor/Log.hs new file mode 100644 --- /dev/null +++ b/src/Mcmc/Monitor/Log.hs @@ -0,0 +1,25 @@ +-- | +-- Module : Mcmc.Monitor.Log +-- Description : Monitor logarithmic values +-- Copyright : (c) Dominik Schrempf, 2020 +-- License : GPL-3.0-or-later +-- +-- Maintainer : dominik.schrempf@gmail.com +-- Stability : unstable +-- Portability : portable +-- +-- Creation date: Fri May 29 12:30:03 2020. +module Mcmc.Monitor.Log + ( renderLog, + ) +where + +import Data.Text.Lazy (Text) +import qualified Data.Text.Lazy.Builder as T +import qualified Data.Text.Lazy.Builder.RealFloat as T +import Numeric.Log + +-- | Print a log value. +renderLog :: Log Double -> Text +renderLog = T.toLazyText . T.formatRealFloat T.Fixed (Just 8) . ln +{-# INLINEABLE renderLog #-} diff --git a/src/Mcmc/Monitor/Parameter.hs b/src/Mcmc/Monitor/Parameter.hs new file mode 100644 --- /dev/null +++ b/src/Mcmc/Monitor/Parameter.hs @@ -0,0 +1,85 @@ +{-# LANGUAGE RankNTypes #-} + +-- | +-- Module : Mcmc.Monitor.Parameter +-- Description : Monitor parameters +-- Copyright : (c) Dominik Schrempf, 2020 +-- License : GPL-3.0-or-later +-- +-- Maintainer : dominik.schrempf@gmail.com +-- Stability : unstable +-- Portability : portable +-- +-- Creation date: Fri May 29 11:11:49 2020. +module Mcmc.Monitor.Parameter + ( MonitorParameter (..), + monitorInt, + monitorRealFloat, + monitorRealFloatF, + monitorRealFloatS, + ) +where + +import Data.Text.Lazy.Builder (Builder) +import qualified Data.Text.Lazy.Builder.Int as T +import qualified Data.Text.Lazy.Builder.RealFloat as T +import Lens.Micro + +-- | Instruction about a parameter to monitor. +data MonitorParameter a + = MonitorParameter + { -- | Name of parameter. + mpName :: String, + -- | Instruction about how to extract the parameter from the state. + mpFunc :: a -> Builder + } + +-- | Monitor integral parameters such as 'Int'. +monitorInt :: + Integral b => + -- | Name of monitor. + String -> + -- | Instruction about which parameter to monitor. + Lens' a b -> + MonitorParameter a +monitorInt n l = MonitorParameter n (\x -> T.decimal $ x ^. l) +{-# SPECIALIZE monitorInt :: String -> Lens' a Int -> MonitorParameter a #-} + +-- | Monitor real float parameters such as 'Double' with eight decimal places +-- (half precision). +monitorRealFloat :: + RealFloat b => + -- | Name of monitor. + String -> + -- | Instruction about which parameter to monitor. + Lens' a b -> + MonitorParameter a +monitorRealFloat n l = + MonitorParameter n (\x -> T.formatRealFloat T.Fixed (Just 8) $ x ^. l) +{-# SPECIALIZE monitorRealFloat :: String -> Lens' a Double -> MonitorParameter a #-} + +-- | Monitor real float parameters such as 'Double' with full precision. +monitorRealFloatF :: + RealFloat b => + -- | Name of monitor. + String -> + -- | Instruction about which parameter to monitor. + Lens' a b -> + MonitorParameter a +monitorRealFloatF n l = MonitorParameter n (\x -> T.realFloat $ x ^. l) +{-# SPECIALIZE monitorRealFloatF :: String -> Lens' a Double -> MonitorParameter a #-} + +-- | Monitor real float parameters such as 'Double' with scientific notation and +-- eight decimal places. +monitorRealFloatS :: + RealFloat b => + -- | Name of monitor. + String -> + -- | Instruction about which parameter to monitor. + Lens' a b -> + MonitorParameter a +monitorRealFloatS n l = + MonitorParameter + n + (\x -> T.formatRealFloat T.Exponent (Just 8) $ x ^. l) +{-# SPECIALIZE monitorRealFloatS :: String -> Lens' a Double -> MonitorParameter a #-} diff --git a/src/Mcmc/Monitor/ParameterBatch.hs b/src/Mcmc/Monitor/ParameterBatch.hs new file mode 100644 --- /dev/null +++ b/src/Mcmc/Monitor/ParameterBatch.hs @@ -0,0 +1,136 @@ +{-# LANGUAGE RankNTypes #-} + +-- | +-- Module : Mcmc.Monitor.ParameterBatch +-- Description : Monitor parameters +-- Copyright : (c) Dominik Schrempf, 2020 +-- License : GPL-3.0-or-later +-- +-- Maintainer : dominik.schrempf@gmail.com +-- Stability : unstable +-- Portability : portable +-- +-- Creation date: Fri May 29 11:11:49 2020. +-- +-- Batch mean monitors. +module Mcmc.Monitor.ParameterBatch + ( MonitorParameterBatch (..), + monitorBatchMeanInt, + monitorBatchMeanIntF, + monitorBatchMeanRealFloat, + monitorBatchMeanRealFloatF, + monitorBatchMeanRealFloatS, + monitorBatchCustom, + ) +where + +import Data.Text.Lazy.Builder (Builder) +import qualified Data.Text.Lazy.Builder.RealFloat as T +import Lens.Micro + +-- | Instruction about a parameter to monitor via batch means. Usually, the +-- monitored parameter is average over the batch size. However, arbitrary +-- functions performing more complicated analyses on the states in the batch can +-- be provided. +-- +-- XXX: Batch monitors are slow at the moment because the monitored parameter +-- has to be extracted from the state for each iteration. +data MonitorParameterBatch a + = MonitorParameterBatch + { -- | Name of batch monitored parameter. + mbpName :: String, + -- | Instruction about how to extract the batch mean from the trace. + mbpFunc :: [a] -> Builder + } + +mapL :: Lens' a b -> [a] -> [b] +mapL l = map (^. l) + +mean :: Real a => [a] -> Double +mean xs = realToFrac (sum xs) / fromIntegral (length xs) +{-# SPECIALIZE mean :: [Double] -> Double #-} +{-# SPECIALIZE mean :: [Int] -> Double #-} + +-- | Batch monitor integral parameters such as 'Int'. Print the mean with eight +-- decimal places (half precision). +monitorBatchMeanInt :: + Integral b => + -- | Name of monitor. + String -> + -- | Instruction about which parameter to monitor. + Lens' a b -> + MonitorParameterBatch a +monitorBatchMeanInt n l = + MonitorParameterBatch + n + (T.formatRealFloat T.Fixed (Just 8) . mean . mapL l) +{-# SPECIALIZE monitorBatchMeanInt :: String -> Lens' a Int -> MonitorParameterBatch a #-} + +-- | Batch monitor integral parameters such as 'Int'. Print the mean with full +-- precision. +monitorBatchMeanIntF :: + Integral b => + -- | Name of monitor. + String -> + -- | Instruction about which parameter to monitor. + Lens' a b -> + MonitorParameterBatch a +monitorBatchMeanIntF n l = + MonitorParameterBatch n (T.realFloat . mean . mapL l) +{-# SPECIALIZE monitorBatchMeanIntF :: String -> Lens' a Int -> MonitorParameterBatch a #-} + +-- | Batch monitor real float parameters such as 'Double' with eight decimal +-- places (half precision). +monitorBatchMeanRealFloat :: + RealFloat b => + -- | Name of monitor. + String -> + -- | Instruction about which parameter to monitor. + Lens' a b -> + MonitorParameterBatch a +monitorBatchMeanRealFloat n l = + MonitorParameterBatch + n + (T.formatRealFloat T.Fixed (Just 8) . mean . mapL l) +{-# SPECIALIZE monitorBatchMeanRealFloat :: String -> Lens' a Double -> MonitorParameterBatch a #-} + +-- | Batch monitor real float parameters such as 'Double' with full precision. +monitorBatchMeanRealFloatF :: + RealFloat b => + -- | Name of monitor. + String -> + -- | Instruction about which parameter to monitor. + Lens' a b -> + MonitorParameterBatch a +monitorBatchMeanRealFloatF n l = + MonitorParameterBatch n (T.realFloat . mean . mapL l) +{-# SPECIALIZE monitorBatchMeanRealFloatF :: String -> Lens' a Double -> MonitorParameterBatch a #-} + +-- | Batch monitor real float parameters such as 'Double' with scientific +-- notation and eight decimal places. +monitorBatchMeanRealFloatS :: + RealFloat b => + -- | Name of monitor. + String -> + -- | Instruction about which parameter to monitor. + Lens' a b -> + MonitorParameterBatch a +monitorBatchMeanRealFloatS n l = + MonitorParameterBatch + n + (T.formatRealFloat T.Exponent (Just 8) . mean . mapL l) +{-# SPECIALIZE monitorBatchMeanRealFloatS :: String -> Lens' a Double -> MonitorParameterBatch a #-} + +-- | Batch monitor parameters with custom lens and builder. +monitorBatchCustom :: + -- | Name of monitor. + String -> + -- | Instruction about which parameter to monitor. + Lens' a b -> + -- | Function to calculate the batch mean. + ([b] -> b) -> + -- | Custom builder. + (b -> Builder) -> + MonitorParameterBatch a +monitorBatchCustom n l f b = + MonitorParameterBatch n (b . f . mapL l) diff --git a/src/Mcmc/Monitor/Time.hs b/src/Mcmc/Monitor/Time.hs new file mode 100644 --- /dev/null +++ b/src/Mcmc/Monitor/Time.hs @@ -0,0 +1,38 @@ +{-# LANGUAGE OverloadedStrings #-} + +-- | +-- Module : Mcmc.Monitor.Time +-- Description : Print time related values +-- Copyright : (c) Dominik Schrempf, 2020 +-- License : GPL-3.0-or-later +-- +-- Maintainer : dominik.schrempf@gmail.com +-- Stability : unstable +-- Portability : portable +-- +-- Creation date: Fri May 29 12:36:43 2020. +module Mcmc.Monitor.Time + ( renderDuration, + ) +where + +import qualified Data.Text.Lazy as T +import Data.Text.Lazy (Text) +import qualified Data.Text.Lazy.Builder as T +import qualified Data.Text.Lazy.Builder.Int as T +import Data.Time.Clock + +-- | Adapted from System.ProgressBar.renderDuration of package +-- [terminal-progressbar-0.4.1](https://hackage.haskell.org/package/terminal-progress-bar-0.4.1). +renderDuration :: NominalDiffTime -> Text +renderDuration dt = hTxt <> mTxt <> sTxt + where + hTxt = renderDecimal h <> ":" + mTxt = renderDecimal m <> ":" + sTxt = renderDecimal s + (h, hRem) = ts `quotRem` 3600 + (m, s) = hRem `quotRem` 60 + -- Total amount of seconds + ts :: Int + ts = round dt + renderDecimal n = T.justifyRight 2 '0' $ T.toLazyText $ T.decimal n diff --git a/src/Mcmc/Move.hs b/src/Mcmc/Move.hs new file mode 100644 --- /dev/null +++ b/src/Mcmc/Move.hs @@ -0,0 +1,285 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} + +-- TODO: Allow execution of moves in order of appearance in the cycle. + +-- TODO: Moves and monitors both use lenses and names for what they move and +-- monitor. Should a data structure be used combining the lens and the name, so +-- that things are cohesive? + +-- TODO: Moves on simplices: SimplexElementScale (?). + +-- TODO: Moves on tree branch lengths. +-- - Slide a node on the tree. +-- - Scale a tree. + +-- TODO: Moves on tree topologies. +-- - NNI +-- - Narrow (what is this, see RevBayes) +-- - FNPR (dito) + +-- TODO: Bactrian moves; https://www.ncbi.nlm.nih.gov/pmc/articles/PMC3845170/. +-- +-- slideBactrian +-- +-- scaleBactrian + +-- | +-- Module : Mcmc.Move +-- Description : Moves and cycles +-- Copyright : (c) Dominik Schrempf 2020 +-- License : GPL-3.0-or-later +-- +-- Maintainer : dominik.schrempf@gmail.com +-- Stability : unstable +-- Portability : portable +-- +-- Creation date: Wed May 20 13:42:53 2020. +module Mcmc.Move + ( -- * Move + Move (..), + MoveSimple (..), + Tuner (tParam, tFunc), + tuner, + tune, + autotune, + + -- * Cycle + Cycle (fromCycle), + fromList, + autotuneC, + summarizeCycle, + + -- * Acceptance + Acceptance (..), + emptyA, + pushA, + resetA, + acceptanceRatios, + ) +where + +import Data.Aeson +import Data.Function +import Data.List +import qualified Data.Map.Strict as M +import Data.Map.Strict (Map) +import Data.Maybe +import qualified Data.Text.Lazy as T +import Data.Text.Lazy (Text) +import qualified Data.Text.Lazy.Builder as B +import qualified Data.Text.Lazy.Builder.Int as B +import qualified Data.Text.Lazy.Builder.RealFloat as B +import Numeric.Log hiding (sum) +import System.Random.MWC + +-- | A 'Move' is an instruction about how the Markov chain will traverse the +-- state space @a@. Essentially, it is a probability density conditioned on the +-- current state. +-- +-- A 'Move' may be tuneable in that it contains information about how to enlarge +-- or shrink the step size to tune the acceptance ratio. +data Move a = Move + { -- | Name (no moves with the same name are allowed in a 'Cycle'). + mvName :: String, + -- | The weight determines how often a 'Move' is executed per iteration of + -- the Markov chain. + mvWeight :: Int, + -- | Simple move without tuning information. + mvSimple :: MoveSimple a, + -- | Tuning is disabled if set to 'Nothing'. + mvTune :: Maybe (Tuner a) + } + +instance Show (Move a) where + show m = show $ mvName m + +instance Eq (Move a) where + m == n = mvName m == mvName n + +instance Ord (Move a) where + compare = compare `on` mvName + +-- XXX: One could also use a different type for 'mvSample', so that +-- 'mvDensity' can be avoided. In detail, +-- +-- @ +-- mvSample :: a -> GenIO -> IO (a, Log Double, Log, Double) +-- @ +-- +-- where the densities describe the probability of going there and back. +-- However, we may need more information about the move for other MCMC samplers +-- different from Metropolis-Hastings. + +-- | Simple move without tuning information. +-- +-- In order to calculate the Metropolis-Hastings ratio, we need to know the +-- probability (density) of jumping forth, and the probability (density) of +-- jumping back. +data MoveSimple a = MoveSimple + { -- | Instruction about randomly moving from the current state to a new + -- state, given some source of randomness. + mvSample :: a -> GenIO -> IO a, + -- | The density of going from one state to another. Set to 'Nothing' for + -- symmetric moves. + mvDensity :: Maybe (a -> a -> Log Double) + } + +-- | Tune the acceptance ratio of a 'Move'; see 'tune', or 'autotune'. +data Tuner a = Tuner + { tParam :: Double, + tFunc :: Double -> MoveSimple a + } + +-- | Create a 'Tuner'. The tuning function accepts a tuning parameter, and +-- returns a corresponding 'MoveSimple'. The larger the tuning parameter, the +-- larger the 'Move', and vice versa. +tuner :: (Double -> MoveSimple a) -> Tuner a +tuner = Tuner 1.0 + +-- | Tune a 'Move'. Return 'Nothing' if 'Move' is not tuneable. If the parameter +-- @dt@ is larger than 1.0, the 'Move' is enlarged, if @0 Move a -> Maybe (Move a) +tune dt tm + | dt <= 0 = error $ "tune: Tuning parameter not positive: " <> show dt <> "." + | otherwise = do + (Tuner t f) <- mvTune tm + let t' = t * dt + return $ tm {mvSimple = f t', mvTune = Just $ Tuner t' f} + +ratioOpt :: Double +ratioOpt = 0.44 + +-- | For a given acceptance ratio, auto tune the 'Move'. For now, a 'Move' is +-- enlarged when the acceptance ratio is above 0.44, and shrunk otherwise. +-- Return 'Nothing' if 'Move' is not tuneable. +-- +-- XXX: The desired acceptance ratio 0.44 is optimal for one-dimensional +-- 'Move's; one could also store the affected number of dimensions with the +-- 'Move' and tune towards an acceptance ratio accounting for the number of +-- dimensions. +autotune :: Double -> Move a -> Maybe (Move a) +autotune a = tune (a / ratioOpt) + +-- | In brief, a 'Cycle' is a list of moves. The state of the Markov chain will +-- be logged only after each 'Cycle', and the iteration counter will be +-- increased by one. __Moves must have unique names__, so that they can be +-- identified. +-- +-- 'Move's are replicated according to their weights and executed in random +-- order. That is, they are not executed in the order they appear in the +-- 'Cycle'. However, if a 'Move' has weight @w@, it is executed exactly @w@ +-- times per iteration. +newtype Cycle a = Cycle {fromCycle :: [Move a]} + +-- | Create a 'Cycle' from a list of 'Move's. +fromList :: [Move a] -> Cycle a +fromList [] = + error "fromList: Received an empty list but cannot create an empty Cycle." +fromList xs = + if length (nub nms) == length nms + then Cycle xs + else error "fromList: Moves don't have unique names." + where + nms = map mvName xs + +-- | Tune the 'Move's in the 'Cycle'. Tuning has no effect on 'Move's that +-- cannot be tuned. See 'autotune'. +autotuneC :: Int -> Acceptance (Move a) -> Cycle a -> Cycle a +autotuneC n a = Cycle . map tuneF . fromCycle + where + ars = acceptanceRatios n a + tuneF m = fromMaybe m (autotune (ars M.! m) m) + +renderRow :: Text -> Text -> Text -> Text -> Text +renderRow name weight acceptRatio tuneParam = " " <> nB <> wB <> rB <> tB + where + nB = T.justifyLeft 30 ' ' name + wB = T.justifyRight 8 ' ' weight + rB = T.justifyRight 20 ' ' acceptRatio + tB = T.justifyRight 20 ' ' tuneParam + +moveHeader :: Text +moveHeader = + renderRow "Move name" "Weight" "Acceptance ratio" "Tuning parameter" + +summarizeMove :: Move a -> Maybe Double -> Text +summarizeMove m r = renderRow (T.pack name) weight acceptRatio tuneParamStr + where + name = mvName m + weight = B.toLazyText $ B.decimal $ mvWeight m + acceptRatio = B.toLazyText $ maybe "" (B.formatRealFloat B.Fixed (Just 3)) r + tuneParamStr = B.toLazyText $ maybe "" (B.formatRealFloat B.Fixed (Just 3)) (tParam <$> mvTune m) + +-- | Summarize the 'Move's in the 'Cycle'. Also report acceptance ratios for the +-- given number of last iterations. +summarizeCycle :: Maybe (Int, Acceptance (Move a)) -> Cycle a -> Text +summarizeCycle acc c = + T.unlines $ + [ "-- Summary of move(s) in cycle.", + -- T.replicate (T.length moveHeader) "─", + moveHeader, + " " <> T.replicate (T.length moveHeader - 3) "─" + ] + ++ [summarizeMove m (ar m) | m <- mvs] + ++ [ " " <> T.replicate (T.length moveHeader - 3) "─", + B.toLazyText $ + B.fromLazyText "-- " + <> B.decimal mpi + <> B.fromString " move(s) per iteration." + <> arStr + ] + where + mvs = fromCycle c + mpi = sum $ map mvWeight mvs + arStr = case acc of + Nothing -> "" + Just (n, _) -> + " Acceptance ratio(s) calculated over " <> B.decimal n <> " iterations." + ars = case acc of + Nothing -> M.empty + Just (n, a) -> acceptanceRatios n a + ar m = ars M.!? m + +-- | For each key @k@, store the list of accepted (True) and rejected (False) +-- proposals. For reasons of efficiency, the lists are stored in reverse order; +-- latest first. +newtype Acceptance k = Acceptance {fromAcceptance :: Map k [Bool]} + +instance ToJSONKey k => ToJSON (Acceptance k) where + toJSON (Acceptance m) = toJSON m + toEncoding (Acceptance m) = toEncoding m + +instance (Ord k, FromJSONKey k) => FromJSON (Acceptance k) where + parseJSON v = Acceptance <$> parseJSON v + +-- | In the beginning there was the Word. +-- +-- Initialize an empty storage of accepted/rejected values. +emptyA :: Ord k => [k] -> Acceptance k +emptyA ks = Acceptance $ M.fromList [(k, []) | k <- ks] + +-- | For key @k@, prepend an accepted (True) or rejected (False) proposal. +pushA :: (Ord k, Show k) => k -> Bool -> Acceptance k -> Acceptance k +-- Unsafe; faster. +pushA k v (Acceptance m) = Acceptance $ M.adjust (v :) k m +-- -- Safe; slower. +-- prependA k v (Acceptance m) | k `M.member` m = Acceptance $ M.adjust (v:) k m +-- | otherwise = error msg +-- where msg = "prependA: Can not add acceptance value for key: " <> show k <> "." +{-# INLINEABLE pushA #-} + +-- | Reset acceptance storage. +resetA :: Acceptance k -> Acceptance k +resetA = Acceptance . M.map (const []) . fromAcceptance + +ratio :: Int -> [Bool] -> Double +ratio n xs = fromIntegral (length ts) / fromIntegral n + where + ts = filter (== True) xs + +-- | Acceptance ratios averaged over the given number of last iterations. If +-- less than @n@ iterations are available, only those are used. +acceptanceRatios :: Int -> Acceptance k -> Map k Double +acceptanceRatios n (Acceptance m) = M.map (ratio n . take n) m diff --git a/src/Mcmc/Move/Generic.hs b/src/Mcmc/Move/Generic.hs new file mode 100644 --- /dev/null +++ b/src/Mcmc/Move/Generic.hs @@ -0,0 +1,136 @@ +{-# LANGUAGE RankNTypes #-} + +-- Technically, only a Getter is needed when calculating the density of the move +-- ('densCont', and similar functions). I tried splitting the lens into a getter +-- and a setter. However, speed improvements were marginal, and some times not +-- even measurable. Using a 'Lens'' is just easier, and has no real drawbacks. + +-- | +-- Module : Mcmc.Move.Generic +-- Description : Generic interface to create moves +-- Copyright : (c) Dominik Schrempf 2020 +-- License : GPL-3.0-or-later +-- +-- Maintainer : dominik.schrempf@gmail.com +-- Stability : unstable +-- Portability : portable +-- +-- Creation date: Thu May 14 20:26:27 2020. +module Mcmc.Move.Generic + ( moveGenericContinuous, + moveSymmetricGenericContinuous, + moveGenericDiscrete, + moveSymmetricGenericDiscrete, + ) +where + +import Lens.Micro +import Mcmc.Move +import Numeric.Log +import Statistics.Distribution +import System.Random.MWC + +jumpCont :: + (ContDistr d, ContGen d) => + Lens' a Double -> + d -> + (Double -> Double -> Double) -> + a -> + GenIO -> + IO a +jumpCont l d f x g = do + dx <- genContVar d g + return $ set l ((x ^. l) `f` dx) x +{-# INLINEABLE jumpCont #-} + +densCont :: + (ContDistr d, ContGen d) => + Lens' a Double -> + d -> + (Double -> Double -> Double) -> + a -> + a -> + Log Double +densCont l d fInv x y = Exp $ logDensity d ((y ^. l) `fInv` (x ^. l)) +{-# INLINEABLE densCont #-} + +-- | Generic function to create moves for continuous parameters ('Double'). +moveGenericContinuous :: + (ContDistr d, ContGen d) => + -- | Instruction about which parameter to change. + Lens' a Double -> + -- | Probability distribution + d -> + -- | Forward operator, e.g. (+), so that x + dx = y. + (Double -> Double -> Double) -> + -- | Inverse operator, e.g.,(-), so that y - dx = x. + (Double -> Double -> Double) -> + MoveSimple a +moveGenericContinuous l d f fInv = + MoveSimple (jumpCont l d f) (Just $ densCont l d fInv) + +-- | Generic function to create symmetric moves for continuous parameters ('Double'). +moveSymmetricGenericContinuous :: + (ContDistr d, ContGen d) => + -- | Instruction about which parameter to change. + Lens' a Double -> + -- | Probability distribution + d -> + -- | Forward operator, e.g. (+), so that x + dx = y. + (Double -> Double -> Double) -> + MoveSimple a +moveSymmetricGenericContinuous l d f = + MoveSimple (jumpCont l d f) Nothing + +jumpDiscrete :: + (DiscreteDistr d, DiscreteGen d) => + Lens' a Int -> + d -> + (Int -> Int -> Int) -> + a -> + GenIO -> + IO a +jumpDiscrete l d f x g = do + dx <- genDiscreteVar d g + return $ set l ((x ^. l) `f` dx) x +{-# INLINEABLE jumpDiscrete #-} + +densDiscrete :: + (DiscreteDistr d, DiscreteGen d) => + Lens' a Int -> + d -> + (Int -> Int -> Int) -> + a -> + a -> + Log Double +densDiscrete l d fInv x y = + Exp $ logProbability d ((y ^. l) `fInv` (x ^. l)) +{-# INLINEABLE densDiscrete #-} + +-- | Generic function to create moves for discrete parameters ('Int'). +moveGenericDiscrete :: + (DiscreteDistr d, DiscreteGen d) => + -- | Instruction about which parameter to change. + Lens' a Int -> + -- | Probability distribution. + d -> + -- | Forward operator, e.g. (+), so that x + dx = y. + (Int -> Int -> Int) -> + -- | Inverse operator, e.g.,(-), so that y - dx = x. + (Int -> Int -> Int) -> + MoveSimple a +moveGenericDiscrete l fd f fInv = + MoveSimple (jumpDiscrete l fd f) (Just $ densDiscrete l fd fInv) + +-- | Generic function to create symmetric moves for discrete parameters ('Int'). +moveSymmetricGenericDiscrete :: + (DiscreteDistr d, DiscreteGen d) => + -- | Instruction about which parameter to change. + Lens' a Int -> + -- | Probability distribution. + d -> + -- | Forward operator, e.g. (+), so that x + dx = y. + (Int -> Int -> Int) -> + MoveSimple a +moveSymmetricGenericDiscrete l fd f = + MoveSimple (jumpDiscrete l fd f) Nothing diff --git a/src/Mcmc/Move/Scale.hs b/src/Mcmc/Move/Scale.hs new file mode 100644 --- /dev/null +++ b/src/Mcmc/Move/Scale.hs @@ -0,0 +1,69 @@ +{-# LANGUAGE RankNTypes #-} + +-- | +-- Module : Mcmc.Move.Scale +-- Description : Scaling move with Gamma distribution +-- Copyright : (c) Dominik Schrempf 2020 +-- License : GPL-3.0-or-later +-- +-- Maintainer : dominik.schrempf@gmail.com +-- Stability : unstable +-- Portability : portable +-- +-- Creation date: Thu May 14 21:49:23 2020. +module Mcmc.Move.Scale + ( scale, + scaleUnbiased, + ) +where + +import Lens.Micro +import Mcmc.Move +import Mcmc.Move.Generic +import Statistics.Distribution.Gamma + +-- The actual move with tuning parameter. +scaleSimple :: Lens' a Double -> Double -> Double -> Double -> MoveSimple a +scaleSimple l k th t = moveGenericContinuous l (gammaDistr k (t * th)) (*) (/) + +-- | Multiplicative move with Gamma distributed density. +scale :: + -- | Name. + String -> + -- | Weight. + Int -> + -- | Instruction about which parameter to change. + Lens' a Double -> + -- | Shape. + Double -> + -- | Scale. + Double -> + -- | Enable tuning. + Bool -> + Move a +scale n w l k th True = + Move n w (scaleSimple l k th 1.0) (Just $ tuner $ scaleSimple l k th) +scale n w l k th False = Move n w (scaleSimple l k th 1.0) Nothing + +-- | Multiplicative move with Gamma distributed density. The scale of the Gamma +-- distributions is set to (shape)^{-1}, so that the mean of the Gamma +-- distribution is 1.0. +scaleUnbiased :: + -- | Name. + String -> + -- | Weight. + Int -> + -- | Instruction about which parameter to change. + Lens' a Double -> + -- | Shape. + Double -> + -- | Enable tuning. + Bool -> + Move a +scaleUnbiased n w l k True = + Move + n + w + (scaleSimple l k (1 / k) 1.0) + (Just $ tuner $ scaleSimple l k (1 / k)) +scaleUnbiased n w l k False = Move n w (scaleSimple l k (1 / k) 1.0) Nothing diff --git a/src/Mcmc/Move/Slide.hs b/src/Mcmc/Move/Slide.hs new file mode 100644 --- /dev/null +++ b/src/Mcmc/Move/Slide.hs @@ -0,0 +1,95 @@ +{-# LANGUAGE RankNTypes #-} + +-- | +-- Module : Mcmc.Move.Slide +-- Description : Normally distributed move +-- Copyright : (c) Dominik Schrempf 2020 +-- License : GPL-3.0-or-later +-- +-- Maintainer : dominik.schrempf@gmail.com +-- Stability : unstable +-- Portability : portable +-- +-- Creation date: Wed May 6 10:59:13 2020. +module Mcmc.Move.Slide + ( slide, + slideSymmetric, + slideUniform, + ) +where + +import Lens.Micro +import Mcmc.Move +import Mcmc.Move.Generic +import Statistics.Distribution.Normal +import Statistics.Distribution.Uniform + +-- The actual move with tuning parameter. +slideSimple :: Lens' a Double -> Double -> Double -> Double -> MoveSimple a +slideSimple l m s t = moveGenericContinuous l (normalDistr m (s * t)) (+) (-) + +-- | Additive move with normally distributed density. +slide :: + -- | Name. + String -> + -- | Weight. + Int -> + -- | Instruction about which parameter to change. + Lens' a Double -> + -- | Mean. + Double -> + -- | Standard deviation. + Double -> + -- | Enable tuning. + Bool -> + Move a +slide n w l m s True = + Move n w (slideSimple l m s 1.0) (Just $ tuner (slideSimple l m s)) +slide n w l m s False = Move n w (slideSimple l m s 1.0) Nothing + +-- The actual move with tuning parameter. +slideSymmetricSimple :: Lens' a Double -> Double -> Double -> MoveSimple a +slideSymmetricSimple l s t = moveSymmetricGenericContinuous l (normalDistr 0.0 (s * t)) (+) + +-- | Additive move with normally distributed density with mean zero. This move +-- is very fast, because the Metropolis-Hastings ratio does not include +-- calculation of the forwards and backwards densities. +slideSymmetric :: + -- | Name. + String -> + -- | Weight. + Int -> + -- | Instruction about which parameter to change. + Lens' a Double -> + -- | Standard deviation. + Double -> + -- | Enable tuning. + Bool -> + Move a +slideSymmetric n w l s True = + Move n w (slideSymmetricSimple l s 1.0) (Just $ tuner (slideSymmetricSimple l s)) +slideSymmetric n w l s False = Move n w (slideSymmetricSimple l s 1.0) Nothing + +-- The actual move with tuning parameter. +slideUniformSimple :: Lens' a Double -> Double -> Double -> MoveSimple a +slideUniformSimple l d t = + moveSymmetricGenericContinuous l (uniformDistr (- t * d) (t * d)) (+) + +-- | Additive move with uniformly distributed density. This move is very fast, +-- because the Metropolis-Hastings ratio does not include calculation of the +-- forwards and backwards densities. +slideUniform :: + -- | Name. + String -> + -- | Weight. + Int -> + -- | Instruction about which parameter to change. + Lens' a Double -> + -- | Delta. + Double -> + -- | Enable tuning. + Bool -> + Move a +slideUniform n w l d True = + Move n w (slideUniformSimple l d 1.0) (Just $ tuner (slideUniformSimple l d)) +slideUniform n w l d False = Move n w (slideUniformSimple l d 1.0) Nothing diff --git a/src/Mcmc/Save.hs b/src/Mcmc/Save.hs new file mode 100644 --- /dev/null +++ b/src/Mcmc/Save.hs @@ -0,0 +1,176 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TemplateHaskell #-} + +-- | +-- Module : Mcmc.Save +-- Description : Save the state of a Markov chain +-- Copyright : (c) Dominik Schrempf, 2020 +-- License : GPL-3.0-or-later +-- +-- Maintainer : dominik.schrempf@gmail.com +-- Stability : unstable +-- Portability : portable +-- +-- Creation date: Tue Jun 16 10:18:54 2020. +-- +-- Save and load an MCMC run. It is easy to save and restore the current state and +-- likelihood (or the trace), but it is not feasible to store all the moves and so +-- on, so they have to be provided again when continuing a run. +module Mcmc.Save + ( saveStatus, + loadStatus, + ) +where + +import Codec.Compression.GZip +import Control.Monad +import Data.Aeson +import Data.Aeson.TH +import qualified Data.ByteString.Lazy as B +import Data.List hiding (cycle) +import Data.Map (Map) +import qualified Data.Map as M +import Data.Time.Clock +import Data.Vector.Unboxed (Vector) +import Data.Word +-- TODO: Remove as soon as split mix is used and is available with the +-- statistics package. +import Mcmc.Item +import Mcmc.Monitor +import Mcmc.Move +import Mcmc.Status hiding (save) +import Mcmc.Trace +import Numeric.Log +import System.IO.Unsafe (unsafePerformIO) +import System.Random.MWC +import Prelude hiding (cycle) + +data Save a + = Save + String + (Item a) + Int + (Trace a) + (Acceptance Int) + (Maybe Int) + (Maybe Int) + Int + (Maybe (Int, UTCTime)) + Bool + (Vector Word32) + +$(deriveJSON defaultOptions 'Save) + +mapKeys :: (Ord k1, Ord k2) => [(k1, k2)] -> Map k1 v -> Map k2 v +mapKeys xs m = foldl' insrt M.empty xs + where + insrt m' (k1, k2) = M.insert k2 (m M.! k1) m' + +toSave :: Status a -> Save a +toSave (Status nm it i tr ac br at is st sv g _ _ c _) = + Save + nm + it + i + tr + ac' + br + at + is + st + sv + g' + where + moveToInt = zip (fromCycle c) [0 ..] + ac' = Acceptance $ mapKeys moveToInt (fromAcceptance ac) + -- TODO: Remove as soon as split mix is used and is available with the + -- statistics package. + g' = fromSeed $ unsafePerformIO $ save g + +-- | Save a 'Status' to file. +-- +-- Saved information: +-- - state +-- - iteration +-- - trace +-- - acceptance ratios +-- - generator +-- +-- Important information that cannot be saved and has to be provided again when +-- a chain is restored: +-- - prior function +-- - likelihood function +-- - cycle +-- - monitor +saveStatus :: ToJSON a => FilePath -> Status a -> IO () +saveStatus fn s = B.writeFile fn $ compress $ encode (toSave s) + +-- fromSav prior lh cycle monitor save +fromSave :: + (a -> Log Double) -> + (a -> Log Double) -> + Cycle a -> + Monitor a -> + Save a -> + Status a +fromSave p l c m (Save nm it i tr ac' br at is st sv g') = + Status + nm + it + i + tr + ac + br + at + is + st + sv + g + p + l + c + m + where + intToMove = zip [0 ..] $ fromCycle c + ac = Acceptance $ mapKeys intToMove (fromAcceptance ac') + -- TODO: Remove as soon as split mix is used and is available with the + -- statistics package. + g = unsafePerformIO $ restore $ toSeed g' + +-- | Load a 'Status' from file. +-- Important information that cannot be saved and has to be provided again when +-- a chain is restored: +-- - prior function +-- - likelihood function +-- - cycle +-- - monitor +loadStatus :: + FromJSON a => + (a -> Log Double) -> + (a -> Log Double) -> + Cycle a -> + Monitor a -> + FilePath -> + IO (Status a) +loadStatus p l c m fn = do + res <- eitherDecode . decompress <$> B.readFile fn + let s = case res of + Left err -> error err + Right sv -> fromSave p l c m sv + -- Check if prior and likelihood matches. + let Item x svp svl = item s + -- Recompute and check the prior and likelihood for the last state because the + -- functions may have changed. Of course, we cannot test for the same + -- function, but having the same prior and likelihood at the last state is + -- already a good indicator. + when + (p x /= svp) + ( error + "loadStatus: Provided prior function does not match the saved prior." + ) + when + (l x /= svl) + ( error + "loadStatus: Provided likelihood function does not match the saved likelihood." + ) + return s diff --git a/src/Mcmc/Status.hs b/src/Mcmc/Status.hs new file mode 100644 --- /dev/null +++ b/src/Mcmc/Status.hs @@ -0,0 +1,143 @@ +-- TODO: Add possibility to store supplementary information about the chain. +-- +-- Maybe something like Trace b; and give a function a -> b to extract +-- supplementary info. + +-- TODO: Status tuned exclusively to the Metropolis-Hastings algorithm. We +-- should abstract the algorithm from the chain. For example, +-- +-- @ +-- data Status a b = Status { Chain a; Algorithm a b} +-- @ +-- +-- where a described the state space and b the auxiliary information of the +-- algorithm. This would also solve the above problem, for example in terms of +-- the Hamiltonian algorithm + +-- | +-- Module : Mcmc.Status +-- Description : What is an MCMC? +-- Copyright : (c) Dominik Schrempf 2020 +-- License : GPL-3.0-or-later +-- +-- Maintainer : dominik.schrempf@gmail.com +-- Stability : unstable +-- Portability : portable +-- +-- Creation date: Tue May 5 18:01:15 2020. +module Mcmc.Status + ( Status (..), + status, + noSave, + ) +where + +import Data.Time.Clock +import Mcmc.Item +import Mcmc.Monitor +import Mcmc.Move +import Mcmc.Trace +import Numeric.Log +import System.Random.MWC hiding (save) +import Prelude hiding (cycle) + +-- | The 'Status' contains all information to run an MCMC chain. It is +-- constructed using the function 'status'. +data Status a = Status + { -- Variables saved to disc. + + -- | The name of the MCMC chain; used as file prefix. + name :: String, + -- | The current 'Item' of the chain combines the current state and the + -- current likelihood. + item :: Item a, + -- | The iteration is the number of completed cycles. + iteration :: Int, + -- | The 'Trace' of the Markov chain in reverse order, the most recent + -- 'Item' is at the head of the list. + trace :: Trace a, + -- | For each 'Move', store the list of accepted (True) and rejected (False) + -- proposals; for reasons of efficiency, the list is also stored in reverse + -- order. + acceptance :: Acceptance (Move a), + -- | Number of burn in iterations; deactivate burn in with 'Nothing'. + burnInIterations :: Maybe Int, + -- | Auto tuning period (only during burn in); deactivate auto tuning with + -- 'Nothing'. + autoTuningPeriod :: Maybe Int, + -- | Number of normal iterations excluding burn in. Note that auto tuning + -- only happens during burn in. + iterations :: Int, + -- | Starting time and starting iteration of chain; used to calculate + -- run time and ETA. + start :: Maybe (Int, UTCTime), + -- | Save the chain? Defaults to 'True'. + save :: Bool, + -- | The random number generator. + generator :: GenIO, + -- Auxiliary functions. + + -- | The prior function. The un-normalized posterior is the product of the + -- prior and the likelihood. + priorF :: a -> Log Double, + -- | The likelihood function. The un-normalized posterior is the product of + -- the prior and the likelihood. + likelihoodF :: a -> Log Double, + -- Variables related to the algorithm. + + -- | A set of 'Move's form a 'Cycle'. + cycle :: Cycle a, + -- | A 'Monitor' observing the chain. + monitor :: Monitor a + } + +-- | Initialize the 'Status' of a Markov chain Monte Carlo run. +status :: + -- | Name of the Markov chain; used as file prefix. + String -> + -- | The prior function. + (a -> Log Double) -> + -- | The likelihood function. + (a -> Log Double) -> + -- | A list of 'Move's executed in forward order. The + -- chain will be logged after each cycle. + Cycle a -> + -- | A 'Monitor' observing the chain. + Monitor a -> + -- | The initial state in the state space @a@. + a -> + -- | Number of burn in iterations; deactivate burn in with 'Nothing'. + Maybe Int -> + -- | Auto tuning period (only during burn in); deactivate + -- auto tuning with 'Nothing'. + Maybe Int -> + -- | Number of normal iterations excluding burn in. Note + -- that auto tuning only happens during burn in. + Int -> + -- | A source of randomness. For reproducible runs, make + -- sure to use a generator with the same seed. + GenIO -> + Status a +status n p l c m x mB mT nI g = + Status + n + i + 0 + (singletonT i) + (emptyA $ fromCycle c) + mB + mT + nI + Nothing + True + g + p + l + c + m + where + i = Item x (p x) (l x) + +-- | Do not save the Markov chain at the end. +noSave :: Status a -> Status a +noSave s = s {save = False} diff --git a/src/Mcmc/Tools/Shuffle.hs b/src/Mcmc/Tools/Shuffle.hs new file mode 100644 --- /dev/null +++ b/src/Mcmc/Tools/Shuffle.hs @@ -0,0 +1,54 @@ +-- | +-- Module : Mcmc.Tools.Shuffle +-- Description : Shuffle a list +-- Copyright : (c) Dominik Schrempf 2020 +-- License : GPL-3.0-or-later +-- +-- Maintainer : dominik.schrempf@gmail.com +-- Stability : unstable +-- Portability : portable +-- +-- Creation date: Wed May 20 14:37:09 2020. +-- +-- From https://wiki.haskell.org/Random_shuffle. +module Mcmc.Tools.Shuffle + ( shuffle, + shuffleN, + grabble, + ) +where + +import Control.Monad +import Control.Monad.ST +import qualified Data.Vector as V +import Data.Vector (Vector) +import qualified Data.Vector.Mutable as M +import System.Random.MWC + ( GenIO, + uniformR, + ) + +-- | Shuffle a list. +shuffle :: [a] -> GenIO -> IO [a] +shuffle xs g = head <$> grabble xs 1 (length xs) g + +-- | Shuffle a list @n@ times. +shuffleN :: [a] -> Int -> GenIO -> IO [[a]] +shuffleN xs n = grabble xs n (length xs) + +-- | @grabble xs m n@ is /O(m*n')/, where @n' = min n (length xs)@. Choose @n'@ +-- elements from @xs@, without replacement, and that @m@ times. +grabble :: [a] -> Int -> Int -> GenIO -> IO [[a]] +grabble xs m n gen = do + swapss <- replicateM m $ forM [0 .. min (l - 1) n] $ \i -> do + j <- uniformR (i, l) gen + return (i, j) + return $ map (V.toList . V.take n . swapElems (V.fromList xs)) swapss + where + l = length xs - 1 + +swapElems :: Vector a -> [(Int, Int)] -> Vector a +swapElems xs swaps = runST $ do + mxs <- V.unsafeThaw xs + mapM_ (uncurry $ M.unsafeSwap mxs) swaps + V.unsafeFreeze mxs diff --git a/src/Mcmc/Trace.hs b/src/Mcmc/Trace.hs new file mode 100644 --- /dev/null +++ b/src/Mcmc/Trace.hs @@ -0,0 +1,60 @@ +-- | +-- Module : Mcmc.Trace +-- Description : Trace of a Markov chain +-- Copyright : (c) Dominik Schrempf 2020 +-- License : GPL-3.0-or-later +-- +-- Maintainer : dominik.schrempf@gmail.com +-- Stability : unstable +-- Portability : portable +-- +-- Creation date: Wed May 20 09:11:25 2020. +module Mcmc.Trace + ( Trace, + singletonT, + pushT, + headT, + takeT, + ) +where + +import Data.Aeson +import Mcmc.Item + +-- | A 'Trace' passes through a list of states with associated likelihoods which +-- are called 'Item's. New 'Item's are prepended, and the path of the Markov +-- chain is stored in reversed order. +newtype Trace a = Trace {fromTrace :: [Item a]} + deriving (Show, Read, Eq) + +instance Semigroup (Trace a) where + (Trace l) <> (Trace r) = Trace (l <> r) + +instance Monoid (Trace a) where + mempty = Trace [] + +instance ToJSON a => ToJSON (Trace a) where + toJSON (Trace xs) = toJSON xs + toEncoding (Trace xs) = toEncoding xs + +instance FromJSON a => FromJSON (Trace a) where + parseJSON v = Trace <$> parseJSONList v + +-- $(deriveJSON defaultOptions 'Trace) + +-- | The empty trace. +singletonT :: Item a -> Trace a +singletonT i = Trace [i] + +-- | Prepend an 'Item' to a 'Trace'. +pushT :: Item a -> Trace a -> Trace a +pushT x = Trace . (:) x . fromTrace +{-# INLINEABLE pushT #-} + +-- | Get the most recent item of the trace. +headT :: Trace a -> Item a +headT = head . fromTrace + +-- | Get the N most recent items of the trace. +takeT :: Int -> Trace a -> [Item a] +takeT n = take n . fromTrace diff --git a/test/Mcmc/Move/SlideSpec.hs b/test/Mcmc/Move/SlideSpec.hs new file mode 100644 --- /dev/null +++ b/test/Mcmc/Move/SlideSpec.hs @@ -0,0 +1,34 @@ +-- | +-- Module : Mcmc.Move.SlideSpec +-- Description : Unit tests for Mcmc.Move.SlideSpec +-- Copyright : (c) Dominik Schrempf 2020 +-- License : GPL-3.0-or-later +-- +-- Maintainer : dominik.schrempf@gmail.com +-- Stability : unstable +-- Portability : portable +-- +-- Creation date: Tue May 19 12:07:43 2020. +module Mcmc.Move.SlideSpec + ( spec, + ) +where + +import Data.Maybe +import Mcmc.Move +import Mcmc.Move.Slide +import Test.Hspec +import Test.QuickCheck + +prop_sym :: Eq b => (a -> a -> b) -> a -> a -> Bool +prop_sym f x y = f x y == f y x + +slideSym :: Move Double +slideSym = slide "symmetric" 1 id 0 1.0 False + +spec :: Spec +spec = + describe "slide" + $ it "has a symmetric proposal distribution if mean is 0" + $ property + $ prop_sym (fromJust $ mvDensity . mvSimple $ slideSym) diff --git a/test/Mcmc/SaveSpec.hs b/test/Mcmc/SaveSpec.hs new file mode 100644 --- /dev/null +++ b/test/Mcmc/SaveSpec.hs @@ -0,0 +1,95 @@ +-- | +-- Module : Mcmc.SaveSpec +-- Description : Unit tests for Mcmc.Save +-- Copyright : (c) Dominik Schrempf, 2020 +-- License : GPL-3.0-or-later +-- +-- Maintainer : dominik.schrempf@gmail.com +-- Stability : unstable +-- Portability : portable +-- +-- Creation date: Tue Jun 16 14:32:55 2020. +module Mcmc.SaveSpec + ( spec, + ) +where + +import Mcmc hiding (save) +import Numeric.Log +import Statistics.Distribution hiding + ( mean, + stdDev, + ) +import Statistics.Distribution.Normal +import System.Random.MWC +import Test.Hspec + +trueMean :: Double +trueMean = 5 + +trueStdDev :: Double +trueStdDev = 4 + +lh :: Double -> Log Double +lh = Exp . logDensity (normalDistr trueMean trueStdDev) + +moveCycle :: Cycle Double +moveCycle = + fromList + [ slideSymmetric "small" 5 id 0.1 True, + slideSymmetric "medium" 2 id 1.0 True, + slideSymmetric "large" 2 id 5.0 True, + slide "skewed" 1 id 1.0 4.0 True + ] + +monStd :: MonitorStdOut Double +monStd = monitorStdOut [monitorRealFloat "mu" id] 10 + +mon :: Monitor Double +mon = Monitor monStd [] [] + +nBurn :: Maybe Int +nBurn = Just 20 + +nAutoTune :: Maybe Int +nAutoTune = Just 10 + +nIter :: Int +nIter = 200 + +spec :: Spec +spec = + describe "saveStatus and loadStatus" + $ it "doesn't change the MCMC chain" + $ do + gen <- create + let s = noSave $ status "SaveSpec" (const 1) lh moveCycle mon 0 nBurn nAutoTune nIter gen + saveStatus "SaveSpec.json" s + s' <- loadStatus (const 1) lh moveCycle mon "SaveSpec.json" + r <- mh s + r' <- mh s' + item r `shouldBe` item r' + iteration r `shouldBe` iteration r' + trace r `shouldBe` trace r' + g <- save $ generator r + g' <- save $ generator r' + g `shouldBe` g' + + -- -- TODO: This will only work with a splittable generator because getNCycles + -- -- changes the generator. + -- describe "mhContinue" + -- $ it "mh 200 + mhContinue 200 == mh 400" + -- $ do + -- gen1 <- create + -- let s1 = noSave $ status "SaveSpec" (const 1) likelihood moveCycle mon 0 nBurn nAutoTune 400 gen1 + -- r1 <- mh s1 + -- gen2 <- create + -- let s2 = noSave $ status "SaveSpec" (const 1) likelihood moveCycle mon 0 nBurn nAutoTune 200 gen2 + -- r2' <- mh s2 + -- r2 <- mhContinue 200 r2' + -- item r1 `shouldBe` item r2 + -- iteration r1 `shouldBe` iteration r2 + -- trace r1 `shouldBe` trace r2 + -- g <- save $ generator r1 + -- g' <- save $ generator r2 + -- g `shouldBe` g' diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,1 @@ +{-# OPTIONS_GHC -F -pgmF hspec-discover #-}