{-# 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
  ( -- * Batch parameter monitors
    MonitorParameterBatch (..),
    (>$<),
    (@#),
    monitorBatchMean,
    monitorBatchMeanF,
    monitorBatchMeanE,
    monitorBatchCustom,
  )
where

import qualified Data.ByteString.Builder as BB
import qualified Data.Double.Conversion.ByteString as BC
import Data.Functor.Contravariant

-- | 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.
--
-- Convert a batch monitor from one data type to another with '(>$<)'.
--
-- For example, batch monitor the mean of the first entry of a tuple:
--
-- @
-- mon = fst >$< monitorBatchMean
-- @
--
-- 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.
    MonitorParameterBatch a -> String
mbpName :: String,
    -- | Instruction about how to extract the batch mean from the trace.
    MonitorParameterBatch a -> [a] -> Builder
mbpFunc :: [a] -> BB.Builder
  }

instance Contravariant (MonitorParameterBatch) where
  contramap :: (a -> b) -> MonitorParameterBatch b -> MonitorParameterBatch a
contramap a -> b
f (MonitorParameterBatch String
n [b] -> Builder
m) = String -> ([a] -> Builder) -> MonitorParameterBatch a
forall a. String -> ([a] -> Builder) -> MonitorParameterBatch a
MonitorParameterBatch String
n ([b] -> Builder
m ([b] -> Builder) -> ([a] -> [b]) -> [a] -> Builder
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b) -> [a] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map a -> b
f)

-- | Convert a batch parameter monitor from one data type to another.
--
-- For example, to batch monitor the mean of the first entry of a tuple:
--
-- @
-- mon = fst @# monitorBatchMean
-- @
(@#) :: (b -> a) -> MonitorParameterBatch a -> MonitorParameterBatch b
@# :: (b -> a) -> MonitorParameterBatch a -> MonitorParameterBatch b
(@#) b -> a
f (MonitorParameterBatch String
n [a] -> Builder
m) = String -> ([b] -> Builder) -> MonitorParameterBatch b
forall a. String -> ([a] -> Builder) -> MonitorParameterBatch a
MonitorParameterBatch String
n ([a] -> Builder
m ([a] -> Builder) -> ([b] -> [a]) -> [b] -> Builder
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (b -> a) -> [b] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map b -> a
f)
{-# DEPRECATED (@#) "Superseded by the contravariant instance, use '(>$<)'." #-}

mean :: Real a => [a] -> Double
mean :: [a] -> Double
mean [a]
xs = a -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac ([a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [a]
xs) Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs)
{-# SPECIALIZE mean :: [Double] -> Double #-}
{-# SPECIALIZE mean :: [Int] -> Double #-}

-- | Batch monitor. Print the mean with eight decimal places (half precision).
monitorBatchMean ::
  Real a =>
  -- | Name.
  String ->
  MonitorParameterBatch a
monitorBatchMean :: String -> MonitorParameterBatch a
monitorBatchMean String
n = String -> ([a] -> Builder) -> MonitorParameterBatch a
forall a. String -> ([a] -> Builder) -> MonitorParameterBatch a
MonitorParameterBatch String
n (ByteString -> Builder
BB.byteString (ByteString -> Builder) -> ([a] -> ByteString) -> [a] -> Builder
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Double -> ByteString
BC.toFixed Int
8 (Double -> ByteString) -> ([a] -> Double) -> [a] -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Double
forall a. Real a => [a] -> Double
mean)
{-# SPECIALIZE monitorBatchMean :: String -> MonitorParameterBatch Int #-}
{-# SPECIALIZE monitorBatchMean :: String -> MonitorParameterBatch Double #-}

-- | Batch monitor. Print the mean with full precision computing the shortest
-- string of digits that correctly represent the number.
monitorBatchMeanF ::
  Real a =>
  -- | Name.
  String ->
  MonitorParameterBatch a
monitorBatchMeanF :: String -> MonitorParameterBatch a
monitorBatchMeanF String
n = String -> ([a] -> Builder) -> MonitorParameterBatch a
forall a. String -> ([a] -> Builder) -> MonitorParameterBatch a
MonitorParameterBatch String
n (ByteString -> Builder
BB.byteString (ByteString -> Builder) -> ([a] -> ByteString) -> [a] -> Builder
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> ByteString
BC.toShortest (Double -> ByteString) -> ([a] -> Double) -> [a] -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Double
forall a. Real a => [a] -> Double
mean)
{-# SPECIALIZE monitorBatchMeanF :: String -> MonitorParameterBatch Int #-}
{-# SPECIALIZE monitorBatchMeanF :: String -> MonitorParameterBatch Double #-}

-- | Batch monitor real float parameters such as 'Double' with scientific
-- notation and eight decimal places.
monitorBatchMeanE ::
  Real a =>
  -- | Name.
  String ->
  MonitorParameterBatch a
monitorBatchMeanE :: String -> MonitorParameterBatch a
monitorBatchMeanE String
n = String -> ([a] -> Builder) -> MonitorParameterBatch a
forall a. String -> ([a] -> Builder) -> MonitorParameterBatch a
MonitorParameterBatch String
n (ByteString -> Builder
BB.byteString (ByteString -> Builder) -> ([a] -> ByteString) -> [a] -> Builder
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Double -> ByteString
BC.toExponential Int
8 (Double -> ByteString) -> ([a] -> Double) -> [a] -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Double
forall a. Real a => [a] -> Double
mean)
{-# SPECIALIZE monitorBatchMeanE :: String -> MonitorParameterBatch Int #-}
{-# SPECIALIZE monitorBatchMeanE :: String -> MonitorParameterBatch Double #-}

-- | Batch monitor parameters with custom lens and builder.
monitorBatchCustom ::
  -- | Name.
  String ->
  -- | Function to calculate the batch mean.
  ([a] -> a) ->
  -- | Custom builder.
  (a -> BB.Builder) ->
  MonitorParameterBatch a
monitorBatchCustom :: String -> ([a] -> a) -> (a -> Builder) -> MonitorParameterBatch a
monitorBatchCustom String
n [a] -> a
f a -> Builder
b = String -> ([a] -> Builder) -> MonitorParameterBatch a
forall a. String -> ([a] -> Builder) -> MonitorParameterBatch a
MonitorParameterBatch String
n (a -> Builder
b (a -> Builder) -> ([a] -> a) -> [a] -> Builder
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> a
f)