{-# LANGUAGE OverloadedStrings #-}
module Mcmc.Mcmc
( Mcmc,
mcmcOutT,
mcmcOutS,
mcmcWarnT,
mcmcWarnS,
mcmcInfoT,
mcmcInfoS,
mcmcDebugT,
mcmcDebugS,
mcmcAutotune,
mcmcResetA,
mcmcSummarizeCycle,
mcmcReport,
mcmcMonitorStdOutHeader,
mcmcMonitorExec,
mcmcRun,
)
where
import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.Trans.State
import Data.Aeson
import qualified Data.ByteString.Lazy.Char8 as BL
import Data.Maybe
import Data.Time.Clock
import Data.Time.Format
import Mcmc.Monitor
import Mcmc.Monitor.Time
import Mcmc.Proposal
import Mcmc.Save
import Mcmc.Status hiding (debug)
import Mcmc.Verbosity
import System.Directory
import System.IO
import Prelude hiding (cycle)
type Mcmc a = StateT (Status a) IO
msgPrepare :: Char -> BL.ByteString -> BL.ByteString
msgPrepare c t = BL.cons c $ ": " <> t
mcmcOutT :: BL.ByteString -> Mcmc a ()
mcmcOutT msg = do
h <- fromMaybe (error "mcmcOut: Log handle is missing.") <$> gets logHandle
liftIO $ BL.putStrLn msg >> BL.hPutStrLn h msg
mcmcOutS :: String -> Mcmc a ()
mcmcOutS = mcmcOutT . BL.pack
mcmcWarnA :: Mcmc a () -> Mcmc a ()
mcmcWarnA a = gets verbosity >>= \v -> info v a
mcmcWarnT :: BL.ByteString -> Mcmc a ()
mcmcWarnT = mcmcWarnA . mcmcOutT . msgPrepare 'W'
mcmcWarnS :: String -> Mcmc a ()
mcmcWarnS = mcmcWarnT . BL.pack
mcmcInfoA :: Mcmc a () -> Mcmc a ()
mcmcInfoA a = gets verbosity >>= \v -> info v a
mcmcInfoT :: BL.ByteString -> Mcmc a ()
mcmcInfoT = mcmcInfoA . mcmcOutT . msgPrepare 'I'
mcmcInfoS :: String -> Mcmc a ()
mcmcInfoS = mcmcInfoT . BL.pack
mcmcDebugA :: Mcmc a () -> Mcmc a ()
mcmcDebugA a = gets verbosity >>= \v -> debug v a
mcmcDebugT :: BL.ByteString -> Mcmc a ()
mcmcDebugT = mcmcDebugA . mcmcOutT . msgPrepare 'D'
mcmcDebugS :: String -> Mcmc a ()
mcmcDebugS = mcmcDebugT . BL.pack
mcmcAutotune :: Mcmc a ()
mcmcAutotune = do
mcmcDebugT "Auto tune."
s <- get
let a = acceptance s
c = cycle s
c' = autotuneCycle a c
put $ s {cycle = c'}
mcmcResetA :: Mcmc a ()
mcmcResetA = do
mcmcDebugT "Reset acceptance ratios."
s <- get
let a = acceptance s
put $ s {acceptance = resetA a}
mcmcSummarizeCycle :: Mcmc a BL.ByteString
mcmcSummarizeCycle = do
a <- gets acceptance
c <- gets cycle
return $ summarizeCycle a c
fTime :: FormatTime t => t -> String
fTime = formatTime defaultTimeLocale "%B %-e, %Y, at %H:%M %P, %Z."
mcmcOpenLog :: Mcmc a ()
mcmcOpenLog = do
s <- get
let lfn = name s ++ ".log"
n = iteration s
frc = forceOverwrite s
fe <- liftIO $ doesFileExist lfn
mh <- liftIO $ case verbosity s of
Quiet -> return Nothing
_ ->
Just <$> case (fe, n, frc) of
(False, _, _) -> openFile lfn WriteMode
(True, 0, True) -> openFile lfn WriteMode
(True, 0, False) -> error "mcmcInit: Log file exists; use 'force' to overwrite output files."
(True, _, _) -> openFile lfn AppendMode
put s {logHandle = mh}
mcmcDebugS $ "Log file name: " ++ lfn ++ "."
mcmcDebugT "Log file opened."
mcmcInit :: Mcmc a ()
mcmcInit = do
mcmcOpenLog
s <- get
t <- liftIO getCurrentTime
mcmcInfoS $ "Start time: " <> fTime t
let m = monitor s
n = iteration s
nm = name s
frc = forceOverwrite s
m' <- if n == 0 then liftIO $ mOpen nm frc m else liftIO $ mAppend nm m
put $ s {monitor = m', start = Just (n, t)}
mcmcReport :: ToJSON a => Mcmc a ()
mcmcReport = do
s <- get
let b = burnInIterations s
t = autoTuningPeriod s
n = iterations s
case b of
Just b' -> mcmcInfoS $ "Burn in for " <> show b' <> " iterations."
Nothing -> return ()
case t of
Just t' -> mcmcInfoS $ "Auto tune every " <> show t' <> " iterations (during burn in only)."
Nothing -> return ()
mcmcInfoS $ "Run chain for " <> show n <> " iterations."
mcmcInfoT "Initial state."
mcmcMonitorStdOutHeader
mcmcMonitorExec
mcmcMonitorStdOutHeader :: Mcmc a ()
mcmcMonitorStdOutHeader = do
m <- gets monitor
mcmcInfoA $ mcmcOutT $ mHeader m
mcmcSave :: ToJSON a => Mcmc a ()
mcmcSave = do
s <- get
if save s
then do
mcmcInfoT "Save Markov chain. For long chains, this may take a while."
liftIO $ saveStatus (name s <> ".mcmc") s
mcmcInfoT "Done saving Markov chain."
else mcmcInfoT "Do not save the Markov chain."
mcmcMonitorExec :: ToJSON a => Mcmc a ()
mcmcMonitorExec = do
s <- get
let i = iteration s
j = iterations s + fromMaybe 0 (burnInIterations s)
m = monitor s
(ss, st) = fromMaybe (error "mcmcMonitorExec: Starting state and time not set.") (start s)
tr = trace s
vb = verbosity s
mt <- liftIO $ mExec vb i ss st tr j m
forM_ mt mcmcOutT
mcmcClose :: ToJSON a => Mcmc a ()
mcmcClose = do
s <- get
mcmcSummarizeCycle >>= mcmcInfoT
mcmcInfoT "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
mcmcInfoT $ "Wall clock run time: " <> renderDuration rt <> "."
mcmcInfoS $ "End time: " <> fTime t
case logHandle s of
Just h -> liftIO $ hClose h
Nothing -> return ()
mcmcRun :: ToJSON a => Mcmc a () -> Status a -> IO (Status a)
mcmcRun algorithm = execStateT $ do
mcmcInit
algorithm
mcmcClose