{-# 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 Data.Maybe
import qualified Data.Text.Lazy as T
import Data.Text.Lazy (Text)
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.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 -> Text -> Text
msgPrepare c t = T.cons c $ ": " <> t
mcmcOutT :: Text -> Mcmc a ()
mcmcOutT msg = do
h <- fromMaybe (error "mcmcOut: Log handle is missing.") <$> gets logHandle
liftIO $ T.putStrLn msg >> T.hPutStrLn h msg
mcmcOutS :: String -> Mcmc a ()
mcmcOutS = mcmcOutT . T.pack
mcmcWarnA :: Mcmc a () -> Mcmc a ()
mcmcWarnA a = gets verbosity >>= \v -> info v a
mcmcWarnT :: Text -> Mcmc a ()
mcmcWarnT = mcmcWarnA . mcmcOutT . msgPrepare 'W'
mcmcWarnS :: String -> Mcmc a ()
mcmcWarnS = mcmcWarnT . T.pack
mcmcInfoA :: Mcmc a () -> Mcmc a ()
mcmcInfoA a = gets verbosity >>= \v -> info v a
mcmcInfoT :: Text -> Mcmc a ()
mcmcInfoT = mcmcInfoA . mcmcOutT . msgPrepare 'I'
mcmcInfoS :: String -> Mcmc a ()
mcmcInfoS = mcmcInfoT . T.pack
mcmcDebugA :: Mcmc a () -> Mcmc a ()
mcmcDebugA a = gets verbosity >>= \v -> debug v a
mcmcDebugT :: Text -> Mcmc a ()
mcmcDebugT = mcmcDebugA . mcmcOutT . msgPrepare 'D'
mcmcDebugS :: String -> Mcmc a ()
mcmcDebugS = mcmcDebugT . T.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 Text
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