{-# LANGUAGE OverloadedStrings #-}

-- |
-- Module      :  Mcmc.Monitor
-- Description :  Monitor a Markov chain
-- Copyright   :  (c) Dominik Schrempf, 2021
-- 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 (..),
    Period,
    simpleMonitor,
    MonitorStdOut,
    monitorStdOut,
    msHeader,
    MonitorFile,
    monitorFile,
    BatchSize,
    MonitorBatch,
    monitorBatch,
    getMonitorBatchSize,

    -- * Use monitors
    mOpen,
    mExec,
    mClose,
  )
where

import Control.Monad
import qualified Data.ByteString.Builder as BB
import qualified Data.ByteString.Lazy.Char8 as BL
import Data.Int
import Data.List hiding (sum)
import Data.Time.Clock
import qualified Data.Vector as VB
import Mcmc.Chain.Link
import Mcmc.Chain.Trace
import Mcmc.Internal.ByteString
import Mcmc.Monitor.Log
import Mcmc.Monitor.Parameter
import Mcmc.Monitor.ParameterBatch
import Mcmc.Monitor.Time
import Mcmc.Settings
import Numeric.Log
import System.IO
import Prelude hiding (sum)

-- | A 'Monitor' observing the chain.
--
-- 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.
    Monitor a -> MonitorStdOut a
mStdOut :: MonitorStdOut a,
    -- | Monitors writing to files.
    Monitor a -> [MonitorFile a]
mFiles :: [MonitorFile a],
    -- | Monitors calculating batch means and
    -- writing to files.
    Monitor a -> [MonitorBatch a]
mBatches :: [MonitorBatch a]
  }

-- | Monitor period.
type Period = Int

-- | Do not monitor parameters.
--
-- Monitor prior and likelihood with given period.
simpleMonitor :: Period -> Monitor a
simpleMonitor :: Period -> Monitor a
simpleMonitor Period
p
  | Period
p Period -> Period -> Bool
forall a. Ord a => a -> a -> Bool
< Period
1 = [Char] -> Monitor a
forall a. HasCallStack => [Char] -> a
error [Char]
"simpleMonitor: Monitor period must be 1 or larger."
  | Bool
otherwise =
    MonitorStdOut a -> [MonitorFile a] -> [MonitorBatch a] -> Monitor a
forall a.
MonitorStdOut a -> [MonitorFile a] -> [MonitorBatch a] -> Monitor a
Monitor ([MonitorParameter a] -> Period -> MonitorStdOut a
forall a. [MonitorParameter a] -> Period -> MonitorStdOut a
MonitorStdOut [] Period
p) [] []

-- | Monitor to standard output; constructed with 'monitorStdOut'.
data MonitorStdOut a = MonitorStdOut
  { MonitorStdOut a -> [MonitorParameter a]
msParams :: [MonitorParameter a],
    MonitorStdOut a -> Period
msPeriod :: Period
  }

-- | Monitor to standard output.
monitorStdOut ::
  [MonitorParameter a] ->
  Period ->
  MonitorStdOut a
monitorStdOut :: [MonitorParameter a] -> Period -> MonitorStdOut a
monitorStdOut [MonitorParameter a]
ps Period
p
  | Period
p Period -> Period -> Bool
forall a. Ord a => a -> a -> Bool
< Period
1 = [Char] -> MonitorStdOut a
forall a. HasCallStack => [Char] -> a
error [Char]
"monitorStdOut: Monitor period must be 1 or larger."
  | Bool
otherwise = [MonitorParameter a] -> Period -> MonitorStdOut a
forall a. [MonitorParameter a] -> Period -> MonitorStdOut a
MonitorStdOut [MonitorParameter a]
ps Period
p

msIWidth :: Int
msIWidth :: Period
msIWidth = Period
9

msWidth :: Int
msWidth :: Period
msWidth = Period
22

msRenderRow :: [BL.ByteString] -> BL.ByteString
msRenderRow :: [ByteString] -> ByteString
msRenderRow [ByteString]
xs = Period -> ByteString -> ByteString
alignRight Period
msIWidth ([ByteString] -> ByteString
forall a. [a] -> a
head [ByteString]
xs) ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> [ByteString] -> ByteString
BL.concat [ByteString]
vals
  where
    vals :: [ByteString]
vals = (ByteString -> ByteString) -> [ByteString] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
map (Period -> ByteString -> ByteString
alignRight Period
msWidth) ([ByteString] -> [ByteString]
forall a. [a] -> [a]
tail [ByteString]
xs)

-- | Header of monitor to standard output.
msHeader :: MonitorStdOut a -> BL.ByteString
msHeader :: MonitorStdOut a -> ByteString
msHeader MonitorStdOut a
m = ByteString -> [ByteString] -> ByteString
BL.intercalate ByteString
"\n" [ByteString
row, ByteString
sep]
  where
    row :: ByteString
row =
      [ByteString] -> ByteString
msRenderRow ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$
        [ByteString
"Iteration", ByteString
"Log-Prior", ByteString
"Log-Likelihood", ByteString
"Log-Posterior"]
          [ByteString] -> [ByteString] -> [ByteString]
forall a. [a] -> [a] -> [a]
++ [ByteString]
nms
          [ByteString] -> [ByteString] -> [ByteString]
forall a. [a] -> [a] -> [a]
++ [ByteString
"Runtime", ByteString
"ETA"]
    sep :: ByteString
sep = Int64 -> Char -> ByteString
BL.replicate (ByteString -> Int64
BL.length ByteString
row) Char
'-'
    nms :: [ByteString]
nms = [[Char] -> ByteString
BL.pack ([Char] -> ByteString) -> [Char] -> ByteString
forall a b. (a -> b) -> a -> b
$ MonitorParameter a -> [Char]
forall a. MonitorParameter a -> [Char]
mpName MonitorParameter a
p | MonitorParameter a
p <- MonitorStdOut a -> [MonitorParameter a]
forall a. MonitorStdOut a -> [MonitorParameter a]
msParams MonitorStdOut a
m]

msDataLine ::
  Int ->
  Link a ->
  Int ->
  UTCTime ->
  Int ->
  MonitorStdOut a ->
  IO BL.ByteString
msDataLine :: Period
-> Link a
-> Period
-> UTCTime
-> Period
-> MonitorStdOut a
-> IO ByteString
msDataLine Period
i (Link a
x Prior
p Prior
l) Period
ss UTCTime
st Period
j MonitorStdOut a
m = do
  UTCTime
ct <- IO UTCTime
getCurrentTime
  let dt :: NominalDiffTime
dt = UTCTime
ct UTCTime -> UTCTime -> NominalDiffTime
`diffUTCTime` UTCTime
st
      -- NOTE: Don't evaluate this when i == ss.
      timePerIter :: NominalDiffTime
timePerIter = NominalDiffTime
dt NominalDiffTime -> NominalDiffTime -> NominalDiffTime
forall a. Fractional a => a -> a -> a
/ Period -> NominalDiffTime
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Period
i Period -> Period -> Period
forall a. Num a => a -> a -> a
- Period
ss)
      -- -- Always 0; doesn't make much sense.
      -- tpi = if (i - ss) < 10
      --       then ""
      --       else renderDurationS timePerIter
      eta :: ByteString
eta =
        if (Period
i Period -> Period -> Period
forall a. Num a => a -> a -> a
- Period
ss) Period -> Period -> Bool
forall a. Ord a => a -> a -> Bool
< Period
10
          then ByteString
""
          else NominalDiffTime -> ByteString
renderDuration (NominalDiffTime -> ByteString) -> NominalDiffTime -> ByteString
forall a b. (a -> b) -> a -> b
$ NominalDiffTime
timePerIter NominalDiffTime -> NominalDiffTime -> NominalDiffTime
forall a. Num a => a -> a -> a
* Period -> NominalDiffTime
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Period
j Period -> Period -> Period
forall a. Num a => a -> a -> a
- Period
i)
  ByteString -> IO ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$
    [ByteString] -> ByteString
msRenderRow ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$
      [[Char] -> ByteString
BL.pack (Period -> [Char]
forall a. Show a => a -> [Char]
show Period
i), Prior -> ByteString
renderLog Prior
p, Prior -> ByteString
renderLog Prior
l, Prior -> ByteString
renderLog (Prior
p Prior -> Prior -> Prior
forall a. Num a => a -> a -> a
* Prior
l)]
        [ByteString] -> [ByteString] -> [ByteString]
forall a. [a] -> [a] -> [a]
++ [Builder -> ByteString
BB.toLazyByteString (Builder -> ByteString) -> Builder -> ByteString
forall a b. (a -> b) -> a -> b
$ MonitorParameter a -> a -> Builder
forall a. MonitorParameter a -> a -> Builder
mpFunc MonitorParameter a
mp a
x | MonitorParameter a
mp <- MonitorStdOut a -> [MonitorParameter a]
forall a. MonitorStdOut a -> [MonitorParameter a]
msParams MonitorStdOut a
m]
        [ByteString] -> [ByteString] -> [ByteString]
forall a. [a] -> [a] -> [a]
++ [NominalDiffTime -> ByteString
renderDuration NominalDiffTime
dt, ByteString
eta]

msExec ::
  Int ->
  Link a ->
  Int ->
  UTCTime ->
  Int ->
  MonitorStdOut a ->
  IO (Maybe BL.ByteString)
msExec :: Period
-> Link a
-> Period
-> UTCTime
-> Period
-> MonitorStdOut a
-> IO (Maybe ByteString)
msExec Period
i Link a
it Period
ss UTCTime
st Period
j MonitorStdOut a
m
  | Period
i Period -> Period -> Period
forall a. Integral a => a -> a -> a
`mod` MonitorStdOut a -> Period
forall a. MonitorStdOut a -> Period
msPeriod MonitorStdOut a
m Period -> Period -> Bool
forall a. Eq a => a -> a -> Bool
/= Period
0 = Maybe ByteString -> IO (Maybe ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe ByteString
forall a. Maybe a
Nothing
  -- -- | i `mod` (msPeriod m * 100) == 0 = do
  -- --   l <- msDataLine i it ss st j m
  -- --   return $ Just $ msHeader m <> "\n" <> l
  | Bool
otherwise = ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just (ByteString -> Maybe ByteString)
-> IO ByteString -> IO (Maybe ByteString)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Period
-> Link a
-> Period
-> UTCTime
-> Period
-> MonitorStdOut a
-> IO ByteString
forall a.
Period
-> Link a
-> Period
-> UTCTime
-> Period
-> MonitorStdOut a
-> IO ByteString
msDataLine Period
i Link a
it Period
ss UTCTime
st Period
j MonitorStdOut a
m

-- | Monitor to a file; constructed with 'monitorFile'.
data MonitorFile a = MonitorFile
  { MonitorFile a -> [Char]
mfName :: String,
    MonitorFile a -> Maybe Handle
mfHandle :: Maybe Handle,
    MonitorFile a -> [MonitorParameter a]
mfParams :: [MonitorParameter a],
    MonitorFile a -> Period
mfPeriod :: Period
  }

-- | Monitor parameters to a file.
monitorFile ::
  -- | Name; used as part of the file name.
  String ->
  [MonitorParameter a] ->
  Period ->
  MonitorFile a
monitorFile :: [Char] -> [MonitorParameter a] -> Period -> MonitorFile a
monitorFile [Char]
n [MonitorParameter a]
ps Period
p
  | Period
p Period -> Period -> Bool
forall a. Ord a => a -> a -> Bool
< Period
1 = [Char] -> MonitorFile a
forall a. HasCallStack => [Char] -> a
error [Char]
"monitorFile: Monitor period must be 1 or larger."
  | Bool
otherwise = [Char]
-> Maybe Handle -> [MonitorParameter a] -> Period -> MonitorFile a
forall a.
[Char]
-> Maybe Handle -> [MonitorParameter a] -> Period -> MonitorFile a
MonitorFile [Char]
n Maybe Handle
forall a. Maybe a
Nothing [MonitorParameter a]
ps Period
p

mfRenderRow :: [BL.ByteString] -> BL.ByteString
mfRenderRow :: [ByteString] -> ByteString
mfRenderRow = ByteString -> [ByteString] -> ByteString
BL.intercalate ByteString
"\t"

mfOpen :: String -> String -> ExecutionMode -> MonitorFile a -> IO (MonitorFile a)
mfOpen :: [Char]
-> [Char] -> ExecutionMode -> MonitorFile a -> IO (MonitorFile a)
mfOpen [Char]
pre [Char]
suf ExecutionMode
em MonitorFile a
m = do
  let fn :: [Char]
fn = [Char] -> [[Char]] -> [Char]
forall a. [a] -> [[a]] -> [a]
intercalate [Char]
"." ([[Char]] -> [Char]) -> [[Char]] -> [Char]
forall a b. (a -> b) -> a -> b
$ ([Char] -> Bool) -> [[Char]] -> [[Char]]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> ([Char] -> Bool) -> [Char] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null) [[Char]
pre, MonitorFile a -> [Char]
forall a. MonitorFile a -> [Char]
mfName MonitorFile a
m, [Char]
suf, [Char]
"monitor"]
  Handle
h <- ExecutionMode -> [Char] -> IO Handle
openWithExecutionMode ExecutionMode
em [Char]
fn
  MonitorFile a -> IO (MonitorFile a)
forall (m :: * -> *) a. Monad m => a -> m a
return (MonitorFile a -> IO (MonitorFile a))
-> MonitorFile a -> IO (MonitorFile a)
forall a b. (a -> b) -> a -> b
$ MonitorFile a
m {mfHandle :: Maybe Handle
mfHandle = Handle -> Maybe Handle
forall a. a -> Maybe a
Just Handle
h}

mfHeader :: MonitorFile a -> IO ()
mfHeader :: MonitorFile a -> IO ()
mfHeader MonitorFile a
m = case MonitorFile a -> Maybe Handle
forall a. MonitorFile a -> Maybe Handle
mfHandle MonitorFile a
m of
  Maybe Handle
Nothing ->
    [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO ()) -> [Char] -> IO ()
forall a b. (a -> b) -> a -> b
$
      [Char]
"mfHeader: No handle available for monitor with name "
        [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> MonitorFile a -> [Char]
forall a. MonitorFile a -> [Char]
mfName MonitorFile a
m
        [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
"."
  Just Handle
h ->
    Handle -> ByteString -> IO ()
BL.hPutStrLn Handle
h (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$
      [ByteString] -> ByteString
mfRenderRow ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$
        [ByteString
"Iteration", ByteString
"Log-Prior", ByteString
"Log-Likelihood", ByteString
"Log-Posterior"]
          [ByteString] -> [ByteString] -> [ByteString]
forall a. [a] -> [a] -> [a]
++ [[Char] -> ByteString
BL.pack ([Char] -> ByteString) -> [Char] -> ByteString
forall a b. (a -> b) -> a -> b
$ MonitorParameter a -> [Char]
forall a. MonitorParameter a -> [Char]
mpName MonitorParameter a
p | MonitorParameter a
p <- MonitorFile a -> [MonitorParameter a]
forall a. MonitorFile a -> [MonitorParameter a]
mfParams MonitorFile a
m]

mfExec ::
  Int ->
  Link a ->
  MonitorFile a ->
  IO ()
mfExec :: Period -> Link a -> MonitorFile a -> IO ()
mfExec Period
i (Link a
x Prior
p Prior
l) MonitorFile a
m
  | Period
i Period -> Period -> Period
forall a. Integral a => a -> a -> a
`mod` MonitorFile a -> Period
forall a. MonitorFile a -> Period
mfPeriod MonitorFile a
m Period -> Period -> Bool
forall a. Eq a => a -> a -> Bool
/= Period
0 = () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
  | Bool
otherwise = case MonitorFile a -> Maybe Handle
forall a. MonitorFile a -> Maybe Handle
mfHandle MonitorFile a
m of
    Maybe Handle
Nothing ->
      [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO ()) -> [Char] -> IO ()
forall a b. (a -> b) -> a -> b
$
        [Char]
"mfExec: No handle available for monitor with name "
          [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> MonitorFile a -> [Char]
forall a. MonitorFile a -> [Char]
mfName MonitorFile a
m
          [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
"."
    Just Handle
h ->
      Handle -> ByteString -> IO ()
BL.hPutStrLn Handle
h (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$
        [ByteString] -> ByteString
mfRenderRow ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$
          [Char] -> ByteString
BL.pack (Period -> [Char]
forall a. Show a => a -> [Char]
show Period
i) ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
:
          Prior -> ByteString
renderLog Prior
p ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
:
          Prior -> ByteString
renderLog Prior
l ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
:
          Prior -> ByteString
renderLog (Prior
p Prior -> Prior -> Prior
forall a. Num a => a -> a -> a
* Prior
l) ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
:
            [Builder -> ByteString
BB.toLazyByteString (Builder -> ByteString) -> Builder -> ByteString
forall a b. (a -> b) -> a -> b
$ MonitorParameter a -> a -> Builder
forall a. MonitorParameter a -> a -> Builder
mpFunc MonitorParameter a
mp a
x | MonitorParameter a
mp <- MonitorFile a -> [MonitorParameter a]
forall a. MonitorFile a -> [MonitorParameter a]
mfParams MonitorFile a
m]

mfClose :: MonitorFile a -> IO ()
mfClose :: MonitorFile a -> IO ()
mfClose MonitorFile a
m = case MonitorFile a -> Maybe Handle
forall a. MonitorFile a -> Maybe Handle
mfHandle MonitorFile a
m of
  Just Handle
h -> Handle -> IO ()
hClose Handle
h
  Maybe Handle
Nothing -> [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO ()) -> [Char] -> IO ()
forall a b. (a -> b) -> a -> b
$ [Char]
"mfClose: File was not opened for monitor " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> MonitorFile a -> [Char]
forall a. MonitorFile a -> [Char]
mfName MonitorFile a
m [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
"."

-- | Batch size.
type BatchSize = Int

-- | Batch monitor to a file.
--
-- Calculate summary statistics over the last given number of iterations (batch
-- size). Construct with 'monitorBatch'.
data MonitorBatch a = MonitorBatch
  { MonitorBatch a -> [Char]
mbName :: String,
    MonitorBatch a -> Maybe Handle
mbHandle :: Maybe Handle,
    MonitorBatch a -> [MonitorParameterBatch a]
mbParams :: [MonitorParameterBatch a],
    MonitorBatch a -> Period
mbSize :: BatchSize
  }

-- | Batch monitor parameters to a file, see 'MonitorBatch'.
monitorBatch ::
  -- | Name; used as part of the file name.
  String ->
  [MonitorParameterBatch a] ->
  BatchSize ->
  MonitorBatch a
monitorBatch :: [Char] -> [MonitorParameterBatch a] -> Period -> MonitorBatch a
monitorBatch [Char]
n [MonitorParameterBatch a]
ps Period
b
  | Period
b Period -> Period -> Bool
forall a. Ord a => a -> a -> Bool
< Period
2 = [Char] -> MonitorBatch a
forall a. HasCallStack => [Char] -> a
error [Char]
"monitorBatch: Batch size must be 2 or larger."
  | Bool
otherwise = [Char]
-> Maybe Handle
-> [MonitorParameterBatch a]
-> Period
-> MonitorBatch a
forall a.
[Char]
-> Maybe Handle
-> [MonitorParameterBatch a]
-> Period
-> MonitorBatch a
MonitorBatch [Char]
n Maybe Handle
forall a. Maybe a
Nothing [MonitorParameterBatch a]
ps Period
b

-- | Batch monitor size.
--
-- Useful to determine the trace length.
getMonitorBatchSize :: MonitorBatch a -> BatchSize
getMonitorBatchSize :: MonitorBatch a -> Period
getMonitorBatchSize = MonitorBatch a -> Period
forall a. MonitorBatch a -> Period
mbSize

mbOpen :: String -> String -> ExecutionMode -> MonitorBatch a -> IO (MonitorBatch a)
mbOpen :: [Char]
-> [Char] -> ExecutionMode -> MonitorBatch a -> IO (MonitorBatch a)
mbOpen [Char]
pre [Char]
suf ExecutionMode
em MonitorBatch a
m = do
  let fn :: [Char]
fn = [Char] -> [[Char]] -> [Char]
forall a. [a] -> [[a]] -> [a]
intercalate [Char]
"." ([[Char]] -> [Char]) -> [[Char]] -> [Char]
forall a b. (a -> b) -> a -> b
$ ([Char] -> Bool) -> [[Char]] -> [[Char]]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> ([Char] -> Bool) -> [Char] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null) [[Char]
pre, MonitorBatch a -> [Char]
forall a. MonitorBatch a -> [Char]
mbName MonitorBatch a
m, [Char]
suf, [Char]
"batch"]
  Handle
h <- ExecutionMode -> [Char] -> IO Handle
openWithExecutionMode ExecutionMode
em [Char]
fn
  MonitorBatch a -> IO (MonitorBatch a)
forall (m :: * -> *) a. Monad m => a -> m a
return (MonitorBatch a -> IO (MonitorBatch a))
-> MonitorBatch a -> IO (MonitorBatch a)
forall a b. (a -> b) -> a -> b
$ MonitorBatch a
m {mbHandle :: Maybe Handle
mbHandle = Handle -> Maybe Handle
forall a. a -> Maybe a
Just Handle
h}

mbHeader :: MonitorBatch a -> IO ()
mbHeader :: MonitorBatch a -> IO ()
mbHeader MonitorBatch a
m = case MonitorBatch a -> Maybe Handle
forall a. MonitorBatch a -> Maybe Handle
mbHandle MonitorBatch a
m of
  Maybe Handle
Nothing ->
    [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO ()) -> [Char] -> IO ()
forall a b. (a -> b) -> a -> b
$
      [Char]
"mbHeader: No handle available for batch monitor with name "
        [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> MonitorBatch a -> [Char]
forall a. MonitorBatch a -> [Char]
mbName MonitorBatch a
m
        [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
"."
  Just Handle
h ->
    Handle -> ByteString -> IO ()
BL.hPutStrLn Handle
h (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$
      [ByteString] -> ByteString
mfRenderRow ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$
        [ByteString
"Iteration", ByteString
"Mean log-Prior", ByteString
"Mean log-Likelihood", ByteString
"Mean log-Posterior"]
          [ByteString] -> [ByteString] -> [ByteString]
forall a. [a] -> [a] -> [a]
++ [[Char] -> ByteString
BL.pack ([Char] -> ByteString) -> [Char] -> ByteString
forall a b. (a -> b) -> a -> b
$ MonitorParameterBatch a -> [Char]
forall a. MonitorParameterBatch a -> [Char]
mbpName MonitorParameterBatch a
mbp | MonitorParameterBatch a
mbp <- MonitorBatch a -> [MonitorParameterBatch a]
forall a. MonitorBatch a -> [MonitorParameterBatch a]
mbParams MonitorBatch a
m]

mean :: VB.Vector (Log Double) -> Log Double
mean :: Vector Prior -> Prior
mean Vector Prior
xs = Vector Prior -> Prior
forall a. Num a => Vector a -> a
VB.sum Vector Prior
xs Prior -> Prior -> Prior
forall a. Fractional a => a -> a -> a
/ Period -> Prior
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Vector Prior -> Period
forall a. Vector a -> Period
VB.length Vector Prior
xs)

mbExec ::
  Int ->
  Trace a ->
  MonitorBatch a ->
  IO ()
mbExec :: Period -> Trace a -> MonitorBatch a -> IO ()
mbExec Period
i Trace a
t MonitorBatch a
m
  | (Period
i Period -> Period -> Period
forall a. Integral a => a -> a -> a
`mod` MonitorBatch a -> Period
forall a. MonitorBatch a -> Period
mbSize MonitorBatch a
m Period -> Period -> Bool
forall a. Eq a => a -> a -> Bool
/= Period
0) Bool -> Bool -> Bool
|| (Period
i Period -> Period -> Bool
forall a. Eq a => a -> a -> Bool
== Period
0) = () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
  | Bool
otherwise = case MonitorBatch a -> Maybe Handle
forall a. MonitorBatch a -> Maybe Handle
mbHandle MonitorBatch a
m of
    Maybe Handle
Nothing ->
      [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO ()) -> [Char] -> IO ()
forall a b. (a -> b) -> a -> b
$
        [Char]
"mbExec: No handle available for batch monitor with name "
          [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> MonitorBatch a -> [Char]
forall a. MonitorBatch a -> [Char]
mbName MonitorBatch a
m
          [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
"."
    Just Handle
h -> do
      Vector (Link a)
xs <- Period -> Trace a -> IO (Vector (Link a))
forall a. Period -> Trace a -> IO (Vector (Link a))
takeT (MonitorBatch a -> Period
forall a. MonitorBatch a -> Period
mbSize MonitorBatch a
m) Trace a
t
      let lps :: Vector Prior
lps = (Link a -> Prior) -> Vector (Link a) -> Vector Prior
forall a b. (a -> b) -> Vector a -> Vector b
VB.map Link a -> Prior
forall a. Link a -> Prior
prior Vector (Link a)
xs
          lls :: Vector Prior
lls = (Link a -> Prior) -> Vector (Link a) -> Vector Prior
forall a b. (a -> b) -> Vector a -> Vector b
VB.map Link a -> Prior
forall a. Link a -> Prior
likelihood Vector (Link a)
xs
          los :: Vector Prior
los = (Prior -> Prior -> Prior)
-> Vector Prior -> Vector Prior -> Vector Prior
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
VB.zipWith Prior -> Prior -> Prior
forall a. Num a => a -> a -> a
(*) Vector Prior
lps Vector Prior
lls
          mlps :: Prior
mlps = Vector Prior -> Prior
mean Vector Prior
lps
          mlls :: Prior
mlls = Vector Prior -> Prior
mean Vector Prior
lls
          mlos :: Prior
mlos = Vector Prior -> Prior
mean Vector Prior
los
      Handle -> ByteString -> IO ()
BL.hPutStrLn Handle
h (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$
        [ByteString] -> ByteString
mfRenderRow ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$
          [Char] -> ByteString
BL.pack (Period -> [Char]
forall a. Show a => a -> [Char]
show Period
i) ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
:
          Prior -> ByteString
renderLog Prior
mlps ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
:
          Prior -> ByteString
renderLog Prior
mlls ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
:
          Prior -> ByteString
renderLog Prior
mlos ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
:
            [Builder -> ByteString
BB.toLazyByteString (Builder -> ByteString) -> Builder -> ByteString
forall a b. (a -> b) -> a -> b
$ MonitorParameterBatch a -> Vector a -> Builder
forall a. MonitorParameterBatch a -> Vector a -> Builder
mbpFunc MonitorParameterBatch a
mbp ((Link a -> a) -> Vector (Link a) -> Vector a
forall a b. (a -> b) -> Vector a -> Vector b
VB.map Link a -> a
forall a. Link a -> a
state Vector (Link a)
xs) | MonitorParameterBatch a
mbp <- MonitorBatch a -> [MonitorParameterBatch a]
forall a. MonitorBatch a -> [MonitorParameterBatch a]
mbParams MonitorBatch a
m]

mbClose :: MonitorBatch a -> IO ()
mbClose :: MonitorBatch a -> IO ()
mbClose MonitorBatch a
m = case MonitorBatch a -> Maybe Handle
forall a. MonitorBatch a -> Maybe Handle
mbHandle MonitorBatch a
m of
  Just Handle
h -> Handle -> IO ()
hClose Handle
h
  Maybe Handle
Nothing -> [Char] -> IO ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO ()) -> [Char] -> IO ()
forall a b. (a -> b) -> a -> b
$ [Char]
"mfClose: File was not opened for batch monitor: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> MonitorBatch a -> [Char]
forall a. MonitorBatch a -> [Char]
mbName MonitorBatch a
m [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
"."

-- | Open the files associated with the 'Monitor'.
mOpen ::
  -- Base name prefix.
  String ->
  -- Base name suffix.
  String ->
  ExecutionMode ->
  Monitor a ->
  IO (Monitor a)
mOpen :: [Char] -> [Char] -> ExecutionMode -> Monitor a -> IO (Monitor a)
mOpen [Char]
pre [Char]
suf ExecutionMode
em (Monitor MonitorStdOut a
s [MonitorFile a]
fs [MonitorBatch a]
bs) = do
  [MonitorFile a]
fs' <- (MonitorFile a -> IO (MonitorFile a))
-> [MonitorFile a] -> IO [MonitorFile a]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([Char]
-> [Char] -> ExecutionMode -> MonitorFile a -> IO (MonitorFile a)
forall a.
[Char]
-> [Char] -> ExecutionMode -> MonitorFile a -> IO (MonitorFile a)
mfOpen [Char]
pre [Char]
suf ExecutionMode
em) [MonitorFile a]
fs
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (ExecutionMode
em ExecutionMode -> ExecutionMode -> Bool
forall a. Eq a => a -> a -> Bool
== ExecutionMode
Continue) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ (MonitorFile a -> IO ()) -> [MonitorFile a] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ MonitorFile a -> IO ()
forall a. MonitorFile a -> IO ()
mfHeader [MonitorFile a]
fs'
  [MonitorBatch a]
bs' <- (MonitorBatch a -> IO (MonitorBatch a))
-> [MonitorBatch a] -> IO [MonitorBatch a]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ([Char]
-> [Char] -> ExecutionMode -> MonitorBatch a -> IO (MonitorBatch a)
forall a.
[Char]
-> [Char] -> ExecutionMode -> MonitorBatch a -> IO (MonitorBatch a)
mbOpen [Char]
pre [Char]
suf ExecutionMode
em) [MonitorBatch a]
bs
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (ExecutionMode
em ExecutionMode -> ExecutionMode -> Bool
forall a. Eq a => a -> a -> Bool
== ExecutionMode
Continue) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ (MonitorBatch a -> IO ()) -> [MonitorBatch a] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ MonitorBatch a -> IO ()
forall a. MonitorBatch a -> IO ()
mbHeader [MonitorBatch a]
bs'
  Handle -> BufferMode -> IO ()
hSetBuffering Handle
stdout BufferMode
LineBuffering
  Monitor a -> IO (Monitor a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Monitor a -> IO (Monitor a)) -> Monitor a -> IO (Monitor a)
forall a b. (a -> b) -> a -> b
$ MonitorStdOut a -> [MonitorFile a] -> [MonitorBatch a] -> Monitor a
forall a.
MonitorStdOut a -> [MonitorFile a] -> [MonitorBatch a] -> Monitor a
Monitor MonitorStdOut a
s [MonitorFile a]
fs' [MonitorBatch a]
bs'

-- | Execute monitors; print status information to files and return text to be
-- printed to standard output and log file.
mExec ::
  Verbosity ->
  -- | Iteration.
  Int ->
  -- | Starting state.
  Int ->
  -- | Starting time.
  UTCTime ->
  Trace a ->
  -- | Total number of iterations; to calculate ETA.
  Int ->
  Monitor a ->
  IO (Maybe BL.ByteString)
mExec :: Verbosity
-> Period
-> Period
-> UTCTime
-> Trace a
-> Period
-> Monitor a
-> IO (Maybe ByteString)
mExec Verbosity
v Period
i Period
ss UTCTime
st Trace a
xs Period
j (Monitor MonitorStdOut a
s [MonitorFile a]
fs [MonitorBatch a]
bs) = do
  Link a
x <- Trace a -> IO (Link a)
forall a. Trace a -> IO (Link a)
headT Trace a
xs
  (MonitorFile a -> IO ()) -> [MonitorFile a] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Period -> Link a -> MonitorFile a -> IO ()
forall a. Period -> Link a -> MonitorFile a -> IO ()
mfExec Period
i Link a
x) [MonitorFile a]
fs
  -- NOTE: Batch monitors are slow because separate batch monitors will extract
  -- separate immutable stacks from the trace. However, using folds on the
  -- mutable stack only could be an option! But then, we require two polymorphic
  -- types (for the fold).
  (MonitorBatch a -> IO ()) -> [MonitorBatch a] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Period -> Trace a -> MonitorBatch a -> IO ()
forall a. Period -> Trace a -> MonitorBatch a -> IO ()
mbExec Period
i Trace a
xs) [MonitorBatch a]
bs
  if Verbosity
v Verbosity -> Verbosity -> Bool
forall a. Eq a => a -> a -> Bool
== Verbosity
Quiet
    then Maybe ByteString -> IO (Maybe ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe ByteString
forall a. Maybe a
Nothing
    else Period
-> Link a
-> Period
-> UTCTime
-> Period
-> MonitorStdOut a
-> IO (Maybe ByteString)
forall a.
Period
-> Link a
-> Period
-> UTCTime
-> Period
-> MonitorStdOut a
-> IO (Maybe ByteString)
msExec Period
i Link a
x Period
ss UTCTime
st Period
j MonitorStdOut a
s

-- | Close the files associated with the 'Monitor'.
mClose :: Monitor a -> IO (Monitor a)
mClose :: Monitor a -> IO (Monitor a)
mClose m :: Monitor a
m@(Monitor MonitorStdOut a
_ [MonitorFile a]
fms [MonitorBatch a]
bms) = do
  (MonitorFile a -> IO ()) -> [MonitorFile a] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ MonitorFile a -> IO ()
forall a. MonitorFile a -> IO ()
mfClose [MonitorFile a]
fms
  (MonitorBatch a -> IO ()) -> [MonitorBatch a] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ MonitorBatch a -> IO ()
forall a. MonitorBatch a -> IO ()
mbClose [MonitorBatch a]
bms
  let fms' :: [MonitorFile a]
fms' = (MonitorFile a -> MonitorFile a)
-> [MonitorFile a] -> [MonitorFile a]
forall a b. (a -> b) -> [a] -> [b]
map (\MonitorFile a
fm -> MonitorFile a
fm {mfHandle :: Maybe Handle
mfHandle = Maybe Handle
forall a. Maybe a
Nothing}) [MonitorFile a]
fms
  let bms' :: [MonitorBatch a]
bms' = (MonitorBatch a -> MonitorBatch a)
-> [MonitorBatch a] -> [MonitorBatch a]
forall a b. (a -> b) -> [a] -> [b]
map (\MonitorBatch a
bm -> MonitorBatch a
bm {mbHandle :: Maybe Handle
mbHandle = Maybe Handle
forall a. Maybe a
Nothing}) [MonitorBatch a]
bms
  Monitor a -> IO (Monitor a)
forall (m :: * -> *) a. Monad m => a -> m a
return Monitor a
m {mFiles :: [MonitorFile a]
mFiles = [MonitorFile a]
fms', mBatches :: [MonitorBatch a]
mBatches = [MonitorBatch a]
bms'}