{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskell #-}

-- |
-- Module      :  Mcmc.Algorithm.MC3
-- Description :  Metropolis-coupled Markov chain Monte Carlo algorithm
-- Copyright   :  2021 Dominik Schrempf
-- License     :  GPL-3.0-or-later
--
-- Maintainer  :  dominik.schrempf@gmail.com
-- Stability   :  unstable
-- Portability :  portable
--
-- Creation date: Mon Nov 23 15:20:33 2020.
--
-- The Metropolis-coupled Markov chain Monte Carlo ('MC3') algorithm.
--
-- Also known as parallel tempering.
--
-- Like any other parallel MCMC algorithm, the 'MC3' algorithm is essentially an
-- 'Mcmc.Algorithm.MHG.MHG' algorithm on the product space of all parallel
-- chains.
--
-- For example, see
--
-- - Geyer, C. J., Markov chain monte carlo maximum likelihood, Computing
--   Science and Statistics, Proceedings of the 23rd Symposium on the Interface,
--   (1991).
--
-- - Altekar, G., Dwarkadas, S., Huelsenbeck, J. P., & Ronquist, F., Parallel
--   metropolis coupled markov chain monte carlo for bayesian phylogenetic
--   inference, Bioinformatics, 20(3), 407–415 (2004).
module Mcmc.Algorithm.MC3
  ( -- * Definitions
    NChains (..),
    SwapPeriod (..),
    NSwaps (..),
    MC3Settings (..),
    MHGChains,
    ReciprocalTemperatures,

    -- * Metropolis-coupled Markov chain Monte Carlo algorithm
    MC3 (..),
    mc3,
    mc3Save,
    mc3Load,
  )
where

import Codec.Compression.GZip
import Control.Concurrent.Async hiding (link)
import Control.Monad
import Data.Aeson
import Data.Aeson.TH
import qualified Data.ByteString.Builder as BB
import qualified Data.ByteString.Lazy.Char8 as BL
import Data.List
import qualified Data.Map.Strict as M
import Data.Time
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as U
import Data.Word
import Mcmc.Acceptance
import Mcmc.Algorithm
import Mcmc.Algorithm.MHG
import Mcmc.Chain.Chain
import Mcmc.Chain.Link
import Mcmc.Chain.Save
import Mcmc.Chain.Trace
import Mcmc.Cycle
import Mcmc.Internal.Random
import Mcmc.Internal.Shuffle
import Mcmc.Likelihood
import Mcmc.Monitor
import Mcmc.Posterior
import Mcmc.Prior
import Mcmc.Proposal
import Mcmc.Settings
import Numeric.Log hiding (sum)
import System.Random.Stateful
import Text.Printf

-- | Total number of parallel chains.
--
-- Must be two or larger.
newtype NChains = NChains {NChains -> Int
fromNChains :: Int}
  deriving (NChains -> NChains -> Bool
(NChains -> NChains -> Bool)
-> (NChains -> NChains -> Bool) -> Eq NChains
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: NChains -> NChains -> Bool
== :: NChains -> NChains -> Bool
$c/= :: NChains -> NChains -> Bool
/= :: NChains -> NChains -> Bool
Eq, ReadPrec [NChains]
ReadPrec NChains
Int -> ReadS NChains
ReadS [NChains]
(Int -> ReadS NChains)
-> ReadS [NChains]
-> ReadPrec NChains
-> ReadPrec [NChains]
-> Read NChains
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
$creadsPrec :: Int -> ReadS NChains
readsPrec :: Int -> ReadS NChains
$creadList :: ReadS [NChains]
readList :: ReadS [NChains]
$creadPrec :: ReadPrec NChains
readPrec :: ReadPrec NChains
$creadListPrec :: ReadPrec [NChains]
readListPrec :: ReadPrec [NChains]
Read, Int -> NChains -> ShowS
[NChains] -> ShowS
NChains -> [Char]
(Int -> NChains -> ShowS)
-> (NChains -> [Char]) -> ([NChains] -> ShowS) -> Show NChains
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> NChains -> ShowS
showsPrec :: Int -> NChains -> ShowS
$cshow :: NChains -> [Char]
show :: NChains -> [Char]
$cshowList :: [NChains] -> ShowS
showList :: [NChains] -> ShowS
Show)

$(deriveJSON defaultOptions ''NChains)

-- | The period of proposing state swaps between chains.
--
-- Must be one or larger.
newtype SwapPeriod = SwapPeriod {SwapPeriod -> Int
fromSwapPeriod :: Int}
  deriving (SwapPeriod -> SwapPeriod -> Bool
(SwapPeriod -> SwapPeriod -> Bool)
-> (SwapPeriod -> SwapPeriod -> Bool) -> Eq SwapPeriod
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: SwapPeriod -> SwapPeriod -> Bool
== :: SwapPeriod -> SwapPeriod -> Bool
$c/= :: SwapPeriod -> SwapPeriod -> Bool
/= :: SwapPeriod -> SwapPeriod -> Bool
Eq, ReadPrec [SwapPeriod]
ReadPrec SwapPeriod
Int -> ReadS SwapPeriod
ReadS [SwapPeriod]
(Int -> ReadS SwapPeriod)
-> ReadS [SwapPeriod]
-> ReadPrec SwapPeriod
-> ReadPrec [SwapPeriod]
-> Read SwapPeriod
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
$creadsPrec :: Int -> ReadS SwapPeriod
readsPrec :: Int -> ReadS SwapPeriod
$creadList :: ReadS [SwapPeriod]
readList :: ReadS [SwapPeriod]
$creadPrec :: ReadPrec SwapPeriod
readPrec :: ReadPrec SwapPeriod
$creadListPrec :: ReadPrec [SwapPeriod]
readListPrec :: ReadPrec [SwapPeriod]
Read, Int -> SwapPeriod -> ShowS
[SwapPeriod] -> ShowS
SwapPeriod -> [Char]
(Int -> SwapPeriod -> ShowS)
-> (SwapPeriod -> [Char])
-> ([SwapPeriod] -> ShowS)
-> Show SwapPeriod
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> SwapPeriod -> ShowS
showsPrec :: Int -> SwapPeriod -> ShowS
$cshow :: SwapPeriod -> [Char]
show :: SwapPeriod -> [Char]
$cshowList :: [SwapPeriod] -> ShowS
showList :: [SwapPeriod] -> ShowS
Show)

$(deriveJSON defaultOptions ''SwapPeriod)

-- | The number of proposed swaps at each swapping event.
--
-- Must be in @[1, NChains - 1]@.
newtype NSwaps = NSwaps {NSwaps -> Int
fromNSwaps :: Int}
  deriving (NSwaps -> NSwaps -> Bool
(NSwaps -> NSwaps -> Bool)
-> (NSwaps -> NSwaps -> Bool) -> Eq NSwaps
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: NSwaps -> NSwaps -> Bool
== :: NSwaps -> NSwaps -> Bool
$c/= :: NSwaps -> NSwaps -> Bool
/= :: NSwaps -> NSwaps -> Bool
Eq, ReadPrec [NSwaps]
ReadPrec NSwaps
Int -> ReadS NSwaps
ReadS [NSwaps]
(Int -> ReadS NSwaps)
-> ReadS [NSwaps]
-> ReadPrec NSwaps
-> ReadPrec [NSwaps]
-> Read NSwaps
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
$creadsPrec :: Int -> ReadS NSwaps
readsPrec :: Int -> ReadS NSwaps
$creadList :: ReadS [NSwaps]
readList :: ReadS [NSwaps]
$creadPrec :: ReadPrec NSwaps
readPrec :: ReadPrec NSwaps
$creadListPrec :: ReadPrec [NSwaps]
readListPrec :: ReadPrec [NSwaps]
Read, Int -> NSwaps -> ShowS
[NSwaps] -> ShowS
NSwaps -> [Char]
(Int -> NSwaps -> ShowS)
-> (NSwaps -> [Char]) -> ([NSwaps] -> ShowS) -> Show NSwaps
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> NSwaps -> ShowS
showsPrec :: Int -> NSwaps -> ShowS
$cshow :: NSwaps -> [Char]
show :: NSwaps -> [Char]
$cshowList :: [NSwaps] -> ShowS
showList :: [NSwaps] -> ShowS
Show)

$(deriveJSON defaultOptions ''NSwaps)

-- | MC3 settings.
data MC3Settings = MC3Settings
  { -- | The number of chains has to be larger equal two.
    MC3Settings -> NChains
mc3NChains :: NChains,
    MC3Settings -> SwapPeriod
mc3SwapPeriod :: SwapPeriod,
    MC3Settings -> NSwaps
mc3NSwaps :: NSwaps
  }
  deriving (MC3Settings -> MC3Settings -> Bool
(MC3Settings -> MC3Settings -> Bool)
-> (MC3Settings -> MC3Settings -> Bool) -> Eq MC3Settings
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: MC3Settings -> MC3Settings -> Bool
== :: MC3Settings -> MC3Settings -> Bool
$c/= :: MC3Settings -> MC3Settings -> Bool
/= :: MC3Settings -> MC3Settings -> Bool
Eq, ReadPrec [MC3Settings]
ReadPrec MC3Settings
Int -> ReadS MC3Settings
ReadS [MC3Settings]
(Int -> ReadS MC3Settings)
-> ReadS [MC3Settings]
-> ReadPrec MC3Settings
-> ReadPrec [MC3Settings]
-> Read MC3Settings
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
$creadsPrec :: Int -> ReadS MC3Settings
readsPrec :: Int -> ReadS MC3Settings
$creadList :: ReadS [MC3Settings]
readList :: ReadS [MC3Settings]
$creadPrec :: ReadPrec MC3Settings
readPrec :: ReadPrec MC3Settings
$creadListPrec :: ReadPrec [MC3Settings]
readListPrec :: ReadPrec [MC3Settings]
Read, Int -> MC3Settings -> ShowS
[MC3Settings] -> ShowS
MC3Settings -> [Char]
(Int -> MC3Settings -> ShowS)
-> (MC3Settings -> [Char])
-> ([MC3Settings] -> ShowS)
-> Show MC3Settings
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> MC3Settings -> ShowS
showsPrec :: Int -> MC3Settings -> ShowS
$cshow :: MC3Settings -> [Char]
show :: MC3Settings -> [Char]
$cshowList :: [MC3Settings] -> ShowS
showList :: [MC3Settings] -> ShowS
Show)

$(deriveJSON defaultOptions ''MC3Settings)

-- | Vector of MHG chains.
type MHGChains a = V.Vector (MHG a)

-- | Vector of reciprocal temperatures.
type ReciprocalTemperatures = U.Vector Double

data SavedMC3 a = SavedMC3
  { forall a. SavedMC3 a -> MC3Settings
savedMC3Settings :: MC3Settings,
    forall a. SavedMC3 a -> Vector (SavedChain a)
savedMC3Chains :: V.Vector (SavedChain a),
    forall a. SavedMC3 a -> ReciprocalTemperatures
savedMC3ReciprocalTemperatures :: ReciprocalTemperatures,
    forall a. SavedMC3 a -> Int
savedMC3Iteration :: Int,
    forall a. SavedMC3 a -> Acceptances Int
savedMC3SwapAcceptance :: Acceptances Int,
    forall a. SavedMC3 a -> (Word64, Word64)
savedMC3Generator :: (Word64, Word64)
  }
  deriving (SavedMC3 a -> SavedMC3 a -> Bool
(SavedMC3 a -> SavedMC3 a -> Bool)
-> (SavedMC3 a -> SavedMC3 a -> Bool) -> Eq (SavedMC3 a)
forall a. Eq a => SavedMC3 a -> SavedMC3 a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => SavedMC3 a -> SavedMC3 a -> Bool
== :: SavedMC3 a -> SavedMC3 a -> Bool
$c/= :: forall a. Eq a => SavedMC3 a -> SavedMC3 a -> Bool
/= :: SavedMC3 a -> SavedMC3 a -> Bool
Eq, Int -> SavedMC3 a -> ShowS
[SavedMC3 a] -> ShowS
SavedMC3 a -> [Char]
(Int -> SavedMC3 a -> ShowS)
-> (SavedMC3 a -> [Char])
-> ([SavedMC3 a] -> ShowS)
-> Show (SavedMC3 a)
forall a. Show a => Int -> SavedMC3 a -> ShowS
forall a. Show a => [SavedMC3 a] -> ShowS
forall a. Show a => SavedMC3 a -> [Char]
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> SavedMC3 a -> ShowS
showsPrec :: Int -> SavedMC3 a -> ShowS
$cshow :: forall a. Show a => SavedMC3 a -> [Char]
show :: SavedMC3 a -> [Char]
$cshowList :: forall a. Show a => [SavedMC3 a] -> ShowS
showList :: [SavedMC3 a] -> ShowS
Show)

$(deriveJSON defaultOptions ''SavedMC3)

toSavedMC3 ::
  MC3 a ->
  IO (SavedMC3 a)
toSavedMC3 :: forall a. MC3 a -> IO (SavedMC3 a)
toSavedMC3 (MC3 MC3Settings
s MHGChains a
mhgs ReciprocalTemperatures
bs Int
i Acceptances Int
ac IOGenM StdGen
g) = do
  Vector (SavedChain a)
scs <- (MHG a -> IO (SavedChain a))
-> MHGChains a -> IO (Vector (SavedChain a))
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Vector a -> m (Vector b)
V.mapM (Chain a -> IO (SavedChain a)
forall a. Chain a -> IO (SavedChain a)
toSavedChain (Chain a -> IO (SavedChain a))
-> (MHG a -> Chain a) -> MHG a -> IO (SavedChain a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MHG a -> Chain a
forall a. MHG a -> Chain a
fromMHG) MHGChains a
mhgs
  (Word64, Word64)
g' <- IOGenM StdGen -> IO (Word64, Word64)
saveGen IOGenM StdGen
g
  SavedMC3 a -> IO (SavedMC3 a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (SavedMC3 a -> IO (SavedMC3 a)) -> SavedMC3 a -> IO (SavedMC3 a)
forall a b. (a -> b) -> a -> b
$ MC3Settings
-> Vector (SavedChain a)
-> ReciprocalTemperatures
-> Int
-> Acceptances Int
-> (Word64, Word64)
-> SavedMC3 a
forall a.
MC3Settings
-> Vector (SavedChain a)
-> ReciprocalTemperatures
-> Int
-> Acceptances Int
-> (Word64, Word64)
-> SavedMC3 a
SavedMC3 MC3Settings
s Vector (SavedChain a)
scs ReciprocalTemperatures
bs Int
i Acceptances Int
ac (Word64, Word64)
g'

fromSavedMC3 ::
  PriorFunction a ->
  LikelihoodFunction a ->
  Cycle a ->
  Monitor a ->
  SavedMC3 a ->
  IO (MC3 a)
fromSavedMC3 :: forall a.
PriorFunction a
-> PriorFunction a
-> Cycle a
-> Monitor a
-> SavedMC3 a
-> IO (MC3 a)
fromSavedMC3 PriorFunction a
pr PriorFunction a
lh Cycle a
cc Monitor a
mn (SavedMC3 MC3Settings
s Vector (SavedChain a)
scs ReciprocalTemperatures
bs Int
i Acceptances Int
ac (Word64, Word64)
g') = do
  Vector (MHG a)
mhgs <-
    [MHG a] -> Vector (MHG a)
forall a. [a] -> Vector a
V.fromList
      ([MHG a] -> Vector (MHG a)) -> IO [MHG a] -> IO (Vector (MHG a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [IO (MHG a)] -> IO [MHG a]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence
        [ Chain a -> MHG a
forall a. Chain a -> MHG a
MHG (Chain a -> MHG a) -> IO (Chain a) -> IO (MHG a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PriorFunction a
-> PriorFunction a
-> Cycle a
-> Monitor a
-> SavedChain a
-> IO (Chain a)
forall a.
PriorFunction a
-> PriorFunction a
-> Cycle a
-> Monitor a
-> SavedChain a
-> IO (Chain a)
fromSavedChain PriorFunction a
pf PriorFunction a
lf Cycle a
cc Monitor a
mn SavedChain a
sc
          | (SavedChain a
sc, PriorFunction a
pf, PriorFunction a
lf) <- [SavedChain a]
-> [PriorFunction a]
-> [PriorFunction a]
-> [(SavedChain a, PriorFunction a, PriorFunction a)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 (Vector (SavedChain a) -> [SavedChain a]
forall a. Vector a -> [a]
V.toList Vector (SavedChain a)
scs) [PriorFunction a]
prs [PriorFunction a]
lhs
        ]
  IOGenM StdGen
g <- (Word64, Word64) -> IO (IOGenM StdGen)
loadGen (Word64, Word64)
g'
  MC3 a -> IO (MC3 a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (MC3 a -> IO (MC3 a)) -> MC3 a -> IO (MC3 a)
forall a b. (a -> b) -> a -> b
$ MC3Settings
-> Vector (MHG a)
-> ReciprocalTemperatures
-> Int
-> Acceptances Int
-> IOGenM StdGen
-> MC3 a
forall a.
MC3Settings
-> MHGChains a
-> ReciprocalTemperatures
-> Int
-> Acceptances Int
-> IOGenM StdGen
-> MC3 a
MC3 MC3Settings
s Vector (MHG a)
mhgs ReciprocalTemperatures
bs Int
i Acceptances Int
ac IOGenM StdGen
g
  where
    prs :: [PriorFunction a]
prs = (Double -> PriorFunction a) -> [Double] -> [PriorFunction a]
forall a b. (a -> b) -> [a] -> [b]
map (PriorFunction a -> Double -> PriorFunction a
forall a. (a -> Log Double) -> Double -> a -> Log Double
heatFunction PriorFunction a
pr) ([Double] -> [PriorFunction a]) -> [Double] -> [PriorFunction a]
forall a b. (a -> b) -> a -> b
$ ReciprocalTemperatures -> [Double]
forall a. Unbox a => Vector a -> [a]
U.toList ReciprocalTemperatures
bs
    lhs :: [PriorFunction a]
lhs = (Double -> PriorFunction a) -> [Double] -> [PriorFunction a]
forall a b. (a -> b) -> [a] -> [b]
map (PriorFunction a -> Double -> PriorFunction a
forall a. (a -> Log Double) -> Double -> a -> Log Double
heatFunction PriorFunction a
lh) ([Double] -> [PriorFunction a]) -> [Double] -> [PriorFunction a]
forall a b. (a -> b) -> a -> b
$ ReciprocalTemperatures -> [Double]
forall a. Unbox a => Vector a -> [a]
U.toList ReciprocalTemperatures
bs

-- | The MC3 algorithm.
data MC3 a = MC3
  { forall a. MC3 a -> MC3Settings
mc3Settings :: MC3Settings,
    -- | The first chain is the cold chain with temperature 1.0.
    forall a. MC3 a -> MHGChains a
mc3MHGChains :: MHGChains a,
    -- | Vector of reciprocal temperatures.
    forall a. MC3 a -> ReciprocalTemperatures
mc3ReciprocalTemperatures :: ReciprocalTemperatures,
    -- | Current iteration.
    forall a. MC3 a -> Int
mc3Iteration :: Int,
    -- | Number of accepted and rejected swaps.
    forall a. MC3 a -> Acceptances Int
mc3SwapAcceptances :: Acceptances Int,
    forall a. MC3 a -> IOGenM StdGen
mc3Generator :: IOGenM StdGen
  }

instance (ToJSON a) => Algorithm (MC3 a) where
  aName :: MC3 a -> [Char]
aName = [Char] -> MC3 a -> [Char]
forall a b. a -> b -> a
const [Char]
"Metropolis-coupled Markov chain Monte Carlo (MC3)"
  aIteration :: MC3 a -> Int
aIteration = MC3 a -> Int
forall a. MC3 a -> Int
mc3Iteration
  aIsInvalidState :: MC3 a -> Bool
aIsInvalidState = MC3 a -> Bool
forall a. ToJSON a => MC3 a -> Bool
mc3IsInvalidState
  aIterate :: IterationMode -> ParallelizationMode -> MC3 a -> IO (MC3 a)
aIterate = IterationMode -> ParallelizationMode -> MC3 a -> IO (MC3 a)
forall a.
ToJSON a =>
IterationMode -> ParallelizationMode -> MC3 a -> IO (MC3 a)
mc3Iterate
  aAutoTune :: TuningType -> Int -> MC3 a -> IO (MC3 a)
aAutoTune = TuningType -> Int -> MC3 a -> IO (MC3 a)
forall a. ToJSON a => TuningType -> Int -> MC3 a -> IO (MC3 a)
mc3AutoTune
  aResetAcceptance :: ResetAcceptance -> MC3 a -> MC3 a
aResetAcceptance = ResetAcceptance -> MC3 a -> MC3 a
forall a. ToJSON a => ResetAcceptance -> MC3 a -> MC3 a
mc3ResetAcceptance
  aCleanAfterBurnIn :: TraceLength -> MC3 a -> IO (MC3 a)
aCleanAfterBurnIn = TraceLength -> MC3 a -> IO (MC3 a)
forall a. ToJSON a => TraceLength -> MC3 a -> IO (MC3 a)
mc3CleanAfterBurnIn
  aSummarizeCycle :: IterationMode -> MC3 a -> ByteString
aSummarizeCycle = IterationMode -> MC3 a -> ByteString
forall a. ToJSON a => IterationMode -> MC3 a -> ByteString
mc3SummarizeCycle
  aOpenMonitors :: AnalysisName -> ExecutionMode -> MC3 a -> IO (MC3 a)
aOpenMonitors = AnalysisName -> ExecutionMode -> MC3 a -> IO (MC3 a)
forall a.
ToJSON a =>
AnalysisName -> ExecutionMode -> MC3 a -> IO (MC3 a)
mc3OpenMonitors
  aExecuteMonitors :: Verbosity -> UTCTime -> Int -> MC3 a -> IO (Maybe ByteString)
aExecuteMonitors = Verbosity -> UTCTime -> Int -> MC3 a -> IO (Maybe ByteString)
forall a.
ToJSON a =>
Verbosity -> UTCTime -> Int -> MC3 a -> IO (Maybe ByteString)
mc3ExecuteMonitors
  aStdMonitorHeader :: MC3 a -> ByteString
aStdMonitorHeader = MC3 a -> ByteString
forall a. ToJSON a => MC3 a -> ByteString
mc3StdMonitorHeader
  aCloseMonitors :: MC3 a -> IO (MC3 a)
aCloseMonitors = MC3 a -> IO (MC3 a)
forall a. ToJSON a => MC3 a -> IO (MC3 a)
mc3CloseMonitors
  aSave :: AnalysisName -> MC3 a -> IO ()
aSave = AnalysisName -> MC3 a -> IO ()
forall a. ToJSON a => AnalysisName -> MC3 a -> IO ()
mc3Save

heatFunction ::
  -- Cold Function.
  (a -> Log Double) ->
  -- Reciprocal temperature.
  Double ->
  -- The heated prior or likelihood function
  (a -> Log Double)
heatFunction :: forall a. (a -> Log Double) -> Double -> a -> Log Double
heatFunction a -> Log Double
f Double
b
  | Double
b Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
<= Double
0 = [Char] -> a -> Log Double
forall a. HasCallStack => [Char] -> a
error [Char]
"heatFunction: Reciprocal temperature is zero or negative."
  | Double
b Double -> Double -> Bool
forall a. Eq a => a -> a -> Bool
== Double
1.0 = a -> Log Double
f
  | Bool
otherwise = (Log Double -> Log Double -> Log Double
forall a. Floating a => a -> a -> a
** Log Double
b') (Log Double -> Log Double) -> (a -> Log Double) -> a -> Log Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Log Double
f
  where
    b' :: Log Double
b' = Double -> Log Double
forall a. a -> Log a
Exp (Double -> Log Double) -> Double -> Log Double
forall a b. (a -> b) -> a -> b
$ Double -> Double
forall a. Floating a => a -> a
log Double
b

--  The prior and likelihood values of the current link are updated.
--
-- NOTE: The trace is not changed! In particular, the prior and likelihood
-- values are not updated for any link of the trace, and no new link is added to
-- the trace.
setReciprocalTemperature ::
  -- Cold prior function.
  PriorFunction a ->
  -- Cold likelihood function.
  LikelihoodFunction a ->
  -- New reciprocal temperature.
  Double ->
  MHG a ->
  MHG a
setReciprocalTemperature :: forall a.
PriorFunction a -> PriorFunction a -> Double -> MHG a -> MHG a
setReciprocalTemperature PriorFunction a
coldPrf PriorFunction a
coldLhf Double
b MHG a
a =
  Chain a -> MHG a
forall a. Chain a -> MHG a
MHG (Chain a -> MHG a) -> Chain a -> MHG a
forall a b. (a -> b) -> a -> b
$
    Chain a
c
      { priorFunction = prf',
        likelihoodFunction = lhf',
        link = Link x (prf' x) (lhf' x)
      }
  where
    c :: Chain a
c = MHG a -> Chain a
forall a. MHG a -> Chain a
fromMHG MHG a
a
    -- We need twice the amount of computations compared to taking the power
    -- after calculating the posterior (pr x * lh x) ** b'. However, I don't
    -- think this is a serious problem.
    --
    -- To minimize computations, it is key to avoid modification of the
    -- reciprocal temperature for the cold chain.
    prf' :: PriorFunction a
prf' = PriorFunction a -> Double -> PriorFunction a
forall a. (a -> Log Double) -> Double -> a -> Log Double
heatFunction PriorFunction a
coldPrf Double
b
    lhf' :: PriorFunction a
lhf' = PriorFunction a -> Double -> PriorFunction a
forall a. (a -> Log Double) -> Double -> a -> Log Double
heatFunction PriorFunction a
coldLhf Double
b
    x :: a
x = Link a -> a
forall a. Link a -> a
state (Link a -> a) -> Link a -> a
forall a b. (a -> b) -> a -> b
$ Chain a -> Link a
forall a. Chain a -> Link a
link Chain a
c

initMHG ::
  -- Cold prior function.
  PriorFunction a ->
  -- Cold likelihood function.
  LikelihoodFunction a ->
  -- Index of MHG chain.
  Int ->
  -- Reciprocal temperature.
  Double ->
  MHG a ->
  IO (MHG a)
initMHG :: forall a.
PriorFunction a
-> PriorFunction a -> Int -> Double -> MHG a -> IO (MHG a)
initMHG PriorFunction a
prf PriorFunction a
lhf Int
i Double
beta MHG a
a
  | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = [Char] -> IO (MHG a)
forall a. HasCallStack => [Char] -> a
error [Char]
"initMHG: Chain index negative."
  -- Do nothing for the cold chain.
  | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = MHG a -> IO (MHG a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (MHG a -> IO (MHG a)) -> MHG a -> IO (MHG a)
forall a b. (a -> b) -> a -> b
$ Chain a -> MHG a
forall a. Chain a -> MHG a
MHG Chain a
c
  | Bool
otherwise = do
      -- We have to push the current link in the trace, since it is not set by
      -- 'setReciprocalTemperature'. The other links in the trace are still
      -- pointing to the link of the cold chain, but this has no effect.
      Trace a
t' <- Link a -> Trace a -> IO (Trace a)
forall a. Link a -> Trace a -> IO (Trace a)
pushT Link a
l Trace a
t
      MHG a -> IO (MHG a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (MHG a -> IO (MHG a)) -> MHG a -> IO (MHG a)
forall a b. (a -> b) -> a -> b
$ Chain a -> MHG a
forall a. Chain a -> MHG a
MHG (Chain a -> MHG a) -> Chain a -> MHG a
forall a b. (a -> b) -> a -> b
$ Chain a
c {trace = t'}
  where
    a' :: MHG a
a' = PriorFunction a -> PriorFunction a -> Double -> MHG a -> MHG a
forall a.
PriorFunction a -> PriorFunction a -> Double -> MHG a -> MHG a
setReciprocalTemperature PriorFunction a
prf PriorFunction a
lhf Double
beta MHG a
a
    c :: Chain a
c = MHG a -> Chain a
forall a. MHG a -> Chain a
fromMHG MHG a
a'
    l :: Link a
l = Chain a -> Link a
forall a. Chain a -> Link a
link Chain a
c
    t :: Trace a
t = Chain a -> Trace a
forall a. Chain a -> Trace a
trace Chain a
c

-- | Initialize an MC3 algorithm with a given number of chains.
--
-- Call 'error' if:
--
-- - The number of chains is one or lower.
--
-- - The swap period is zero or negative.
mc3 ::
  MC3Settings ->
  Settings ->
  PriorFunction a ->
  LikelihoodFunction a ->
  Cycle a ->
  Monitor a ->
  InitialState a ->
  StdGen ->
  IO (MC3 a)
mc3 :: forall a.
MC3Settings
-> Settings
-> PriorFunction a
-> PriorFunction a
-> Cycle a
-> Monitor a
-> a
-> StdGen
-> IO (MC3 a)
mc3 MC3Settings
sMc3 Settings
s PriorFunction a
pr PriorFunction a
lh Cycle a
cc Monitor a
mn a
i0 StdGen
g
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 = [Char] -> IO (MC3 a)
forall a. HasCallStack => [Char] -> a
error [Char]
"mc3: The number of chains must be two or larger."
  | Int
sp Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 = [Char] -> IO (MC3 a)
forall a. HasCallStack => [Char] -> a
error [Char]
"mc3: The swap period must be strictly positive."
  | Int
sn Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 Bool -> Bool -> Bool
|| Int
sn Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 = [Char] -> IO (MC3 a)
forall a. HasCallStack => [Char] -> a
error [Char]
"mc3: The number of swaps must be in [1, NChains - 1]."
  | Bool
otherwise = do
      -- Split random number generator.
      let gs :: [StdGen]
gs = Int -> [StdGen] -> [StdGen]
forall a. Int -> [a] -> [a]
take (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ([StdGen] -> [StdGen]) -> [StdGen] -> [StdGen]
forall a b. (a -> b) -> a -> b
$ (StdGen -> Maybe (StdGen, StdGen)) -> StdGen -> [StdGen]
forall b a. (b -> Maybe (a, b)) -> b -> [a]
unfoldr ((StdGen, StdGen) -> Maybe (StdGen, StdGen)
forall a. a -> Maybe a
Just ((StdGen, StdGen) -> Maybe (StdGen, StdGen))
-> (StdGen -> (StdGen, StdGen)) -> StdGen -> Maybe (StdGen, StdGen)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StdGen -> (StdGen, StdGen)
forall g. RandomGen g => g -> (g, g)
split) StdGen
g
      -- Prepare MHG chains.
      Vector (MHG a)
cs <- (StdGen -> IO (MHG a)) -> Vector StdGen -> IO (Vector (MHG a))
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Vector a -> m (Vector b)
V.mapM (Settings
-> PriorFunction a
-> PriorFunction a
-> Cycle a
-> Monitor a
-> a
-> StdGen
-> IO (MHG a)
forall a.
Settings
-> PriorFunction a
-> PriorFunction a
-> Cycle a
-> Monitor a
-> a
-> StdGen
-> IO (MHG a)
mhg Settings
s PriorFunction a
pr PriorFunction a
lh Cycle a
cc Monitor a
mn a
i0) ([StdGen] -> Vector StdGen
forall a. [a] -> Vector a
V.fromList ([StdGen] -> Vector StdGen) -> [StdGen] -> Vector StdGen
forall a b. (a -> b) -> a -> b
$ [StdGen] -> [StdGen]
forall a. HasCallStack => [a] -> [a]
tail [StdGen]
gs)
      Vector (MHG a)
hcs <- (Int -> Double -> MHG a -> IO (MHG a))
-> Vector Double -> Vector (MHG a) -> IO (Vector (MHG a))
forall (m :: * -> *) a b c.
Monad m =>
(Int -> a -> b -> m c) -> Vector a -> Vector b -> m (Vector c)
V.izipWithM (PriorFunction a
-> PriorFunction a -> Int -> Double -> MHG a -> IO (MHG a)
forall a.
PriorFunction a
-> PriorFunction a -> Int -> Double -> MHG a -> IO (MHG a)
initMHG PriorFunction a
pr PriorFunction a
lh) (ReciprocalTemperatures -> Vector Double
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
V.convert ReciprocalTemperatures
bs) Vector (MHG a)
cs
      -- Do not reuse the initial generator.
      IOGenM StdGen
gm <- StdGen -> IO (IOGenM StdGen)
forall (m :: * -> *) g. MonadIO m => g -> m (IOGenM g)
newIOGenM (StdGen -> IO (IOGenM StdGen)) -> StdGen -> IO (IOGenM StdGen)
forall a b. (a -> b) -> a -> b
$ [StdGen] -> StdGen
forall a. HasCallStack => [a] -> a
head [StdGen]
gs
      MC3 a -> IO (MC3 a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (MC3 a -> IO (MC3 a)) -> MC3 a -> IO (MC3 a)
forall a b. (a -> b) -> a -> b
$ MC3Settings
-> Vector (MHG a)
-> ReciprocalTemperatures
-> Int
-> Acceptances Int
-> IOGenM StdGen
-> MC3 a
forall a.
MC3Settings
-> MHGChains a
-> ReciprocalTemperatures
-> Int
-> Acceptances Int
-> IOGenM StdGen
-> MC3 a
MC3 MC3Settings
sMc3 Vector (MHG a)
hcs ReciprocalTemperatures
bs Int
0 ([Int] -> Acceptances Int
forall k. Ord k => [k] -> Acceptances k
emptyA [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2]) IOGenM StdGen
gm
  where
    n :: Int
n = NChains -> Int
fromNChains (NChains -> Int) -> NChains -> Int
forall a b. (a -> b) -> a -> b
$ MC3Settings -> NChains
mc3NChains MC3Settings
sMc3
    sp :: Int
sp = SwapPeriod -> Int
fromSwapPeriod (SwapPeriod -> Int) -> SwapPeriod -> Int
forall a b. (a -> b) -> a -> b
$ MC3Settings -> SwapPeriod
mc3SwapPeriod MC3Settings
sMc3
    sn :: Int
sn = NSwaps -> Int
fromNSwaps (NSwaps -> Int) -> NSwaps -> Int
forall a b. (a -> b) -> a -> b
$ MC3Settings -> NSwaps
mc3NSwaps MC3Settings
sMc3
    -- NOTE: The initial choice of reciprocal temperatures is based on a few
    -- tests but otherwise pretty arbitrary.
    --
    -- NOTE: Have to 'take n' elements, because vectors are not as lazy as
    -- lists.
    bs :: ReciprocalTemperatures
bs = [Double] -> ReciprocalTemperatures
forall a. Unbox a => [a] -> Vector a
U.fromList ([Double] -> ReciprocalTemperatures)
-> [Double] -> ReciprocalTemperatures
forall a b. (a -> b) -> a -> b
$ Int -> [Double] -> [Double]
forall a. Int -> [a] -> [a]
take Int
n ([Double] -> [Double]) -> [Double] -> [Double]
forall a b. (a -> b) -> a -> b
$ (Double -> Double) -> Double -> [Double]
forall a. (a -> a) -> a -> [a]
iterate (Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
0.97) Double
1.0

mc3Fn :: AnalysisName -> FilePath
mc3Fn :: AnalysisName -> [Char]
mc3Fn (AnalysisName [Char]
nm) = [Char]
nm [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
".mcmc.mc3"

-- | Save an MC3 algorithm.
mc3Save ::
  (ToJSON a) =>
  AnalysisName ->
  MC3 a ->
  IO ()
mc3Save :: forall a. ToJSON a => AnalysisName -> MC3 a -> IO ()
mc3Save AnalysisName
nm MC3 a
a = do
  SavedMC3 a
savedMC3 <- MC3 a -> IO (SavedMC3 a)
forall a. MC3 a -> IO (SavedMC3 a)
toSavedMC3 MC3 a
a
  [Char] -> ByteString -> IO ()
BL.writeFile (AnalysisName -> [Char]
mc3Fn AnalysisName
nm) (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
compress (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ SavedMC3 a -> ByteString
forall a. ToJSON a => a -> ByteString
encode SavedMC3 a
savedMC3

-- | Load an MC3 algorithm.
--
-- Also create a backup of the save.
--
-- See 'Mcmc.Mcmc.mcmcContinue'.
mc3Load ::
  (FromJSON a) =>
  PriorFunction a ->
  LikelihoodFunction a ->
  Cycle a ->
  Monitor a ->
  AnalysisName ->
  IO (MC3 a)
mc3Load :: forall a.
FromJSON a =>
PriorFunction a
-> PriorFunction a
-> Cycle a
-> Monitor a
-> AnalysisName
-> IO (MC3 a)
mc3Load PriorFunction a
pr PriorFunction a
lh Cycle a
cc Monitor a
mn AnalysisName
nm = do
  Either [Char] (SavedMC3 a)
savedMC3 <- ByteString -> Either [Char] (SavedMC3 a)
forall a. FromJSON a => ByteString -> Either [Char] a
eitherDecode (ByteString -> Either [Char] (SavedMC3 a))
-> (ByteString -> ByteString)
-> ByteString
-> Either [Char] (SavedMC3 a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
decompress (ByteString -> Either [Char] (SavedMC3 a))
-> IO ByteString -> IO (Either [Char] (SavedMC3 a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> IO ByteString
BL.readFile [Char]
fn
  ([Char] -> IO (MC3 a))
-> (SavedMC3 a -> IO (MC3 a))
-> Either [Char] (SavedMC3 a)
-> IO (MC3 a)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either [Char] -> IO (MC3 a)
forall a. HasCallStack => [Char] -> a
error (PriorFunction a
-> PriorFunction a
-> Cycle a
-> Monitor a
-> SavedMC3 a
-> IO (MC3 a)
forall a.
PriorFunction a
-> PriorFunction a
-> Cycle a
-> Monitor a
-> SavedMC3 a
-> IO (MC3 a)
fromSavedMC3 PriorFunction a
pr PriorFunction a
lh Cycle a
cc Monitor a
mn) Either [Char] (SavedMC3 a)
savedMC3
  where
    -- fnBak = mc3Fn $ AnalysisName $ (fromAnalysisName nm ++ ".bak")
    fn :: [Char]
fn = AnalysisName -> [Char]
mc3Fn AnalysisName
nm

-- I call the chains left and right, because it is easy to think about them as
-- being left and right. Of course, the left chain may also have a larger index
-- than the right chain.
swapWith ::
  -- Index i>=0 of left chain.
  Int ->
  -- Index j>=0, j/=i of right chain.
  Int ->
  MHGChains a ->
  (MHGChains a, Posterior)
swapWith :: forall a. Int -> Int -> MHGChains a -> (MHGChains a, Log Double)
swapWith Int
i Int
j MHGChains a
xs
  | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = [Char] -> (MHGChains a, Log Double)
forall a. HasCallStack => [Char] -> a
error [Char]
"swapWith: Left index is negative."
  | Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = [Char] -> (MHGChains a, Log Double)
forall a. HasCallStack => [Char] -> a
error [Char]
"swapWith: Right index is negative."
  | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j = [Char] -> (MHGChains a, Log Double)
forall a. HasCallStack => [Char] -> a
error [Char]
"swapWith: Indices are equal."
  | Bool
otherwise = (MHGChains a
xs', Log Double
q)
  where
    -- Gather information from current chains.
    cl :: Chain a
cl = MHG a -> Chain a
forall a. MHG a -> Chain a
fromMHG (MHG a -> Chain a) -> MHG a -> Chain a
forall a b. (a -> b) -> a -> b
$ MHGChains a
xs MHGChains a -> Int -> MHG a
forall a. Vector a -> Int -> a
V.! Int
i
    cr :: Chain a
cr = MHG a -> Chain a
forall a. MHG a -> Chain a
fromMHG (MHG a -> Chain a) -> MHG a -> Chain a
forall a b. (a -> b) -> a -> b
$ MHGChains a
xs MHGChains a -> Int -> MHG a
forall a. Vector a -> Int -> a
V.! Int
j
    ll :: Link a
ll = Chain a -> Link a
forall a. Chain a -> Link a
link Chain a
cl
    lr :: Link a
lr = Chain a -> Link a
forall a. Chain a -> Link a
link Chain a
cr
    prl :: Log Double
prl = Link a -> Log Double
forall a. Link a -> Log Double
prior Link a
ll
    prr :: Log Double
prr = Link a -> Log Double
forall a. Link a -> Log Double
prior Link a
lr
    lhl :: Log Double
lhl = Link a -> Log Double
forall a. Link a -> Log Double
likelihood Link a
ll
    lhr :: Log Double
lhr = Link a -> Log Double
forall a. Link a -> Log Double
likelihood Link a
lr
    -- Swap the states.
    xl' :: a
xl' = Link a -> a
forall a. Link a -> a
state Link a
lr
    xr' :: a
xr' = Link a -> a
forall a. Link a -> a
state Link a
ll
    -- Compute new priors and likelihoods.
    prl' :: Log Double
prl' = Chain a -> PriorFunction a
forall a. Chain a -> PriorFunction a
priorFunction Chain a
cl a
xl'
    prr' :: Log Double
prr' = Chain a -> PriorFunction a
forall a. Chain a -> PriorFunction a
priorFunction Chain a
cr a
xr'
    lhl' :: Log Double
lhl' = Chain a -> PriorFunction a
forall a. Chain a -> PriorFunction a
likelihoodFunction Chain a
cl a
xl'
    lhr' :: Log Double
lhr' = Chain a -> PriorFunction a
forall a. Chain a -> PriorFunction a
likelihoodFunction Chain a
cr a
xr'
    -- Set the new links and the proposed state.
    ll' :: Link a
ll' = a -> Log Double -> Log Double -> Link a
forall a. a -> Log Double -> Log Double -> Link a
Link a
xl' Log Double
prl' Log Double
lhl'
    lr' :: Link a
lr' = a -> Log Double -> Log Double -> Link a
forall a. a -> Log Double -> Log Double -> Link a
Link a
xr' Log Double
prr' Log Double
lhr'
    cl' :: Chain a
cl' = Chain a
cl {link = ll'}
    cr' :: Chain a
cr' = Chain a
cr {link = lr'}
    xs' :: MHGChains a
xs' = MHGChains a
xs MHGChains a -> [(Int, MHG a)] -> MHGChains a
forall a. Vector a -> [(Int, a)] -> Vector a
V.// [(Int
i, Chain a -> MHG a
forall a. Chain a -> MHG a
MHG Chain a
cl'), (Int
j, Chain a -> MHG a
forall a. Chain a -> MHG a
MHG Chain a
cr')]
    -- Compute the Metropolis ratio.
    nominator :: Log Double
nominator = Log Double
prl' Log Double -> Log Double -> Log Double
forall a. Num a => a -> a -> a
* Log Double
prr' Log Double -> Log Double -> Log Double
forall a. Num a => a -> a -> a
* Log Double
lhl' Log Double -> Log Double -> Log Double
forall a. Num a => a -> a -> a
* Log Double
lhr'
    denominator :: Log Double
denominator = Log Double
prl Log Double -> Log Double -> Log Double
forall a. Num a => a -> a -> a
* Log Double
prr Log Double -> Log Double -> Log Double
forall a. Num a => a -> a -> a
* Log Double
lhl Log Double -> Log Double -> Log Double
forall a. Num a => a -> a -> a
* Log Double
lhr
    q :: Log Double
q = Log Double
nominator Log Double -> Log Double -> Log Double
forall a. Fractional a => a -> a -> a
/ Log Double
denominator

mc3ProposeSwap ::
  MC3 a ->
  -- Index of left chain.
  Int ->
  IO (MC3 a)
mc3ProposeSwap :: forall a. MC3 a -> Int -> IO (MC3 a)
mc3ProposeSwap MC3 a
a Int
i = do
  let cs :: MHGChains a
cs = MC3 a -> MHGChains a
forall a. MC3 a -> MHGChains a
mc3MHGChains MC3 a
a
  -- -- Debug.
  -- prL = prior $ link $ fromMHG $ cs V.! i
  -- prR = prior $ link $ fromMHG $ cs V.! (i+1)
  -- lhL = likelihood $ link $ fromMHG $ cs V.! i
  -- lhR = likelihood $ link $ fromMHG $ cs V.! (i+1)
  -- 1. Sample new state and get the Metropolis ratio.
  let (!MHGChains a
y, !Log Double
r) = Int -> Int -> MHGChains a -> (MHGChains a, Log Double)
forall a. Int -> Int -> MHGChains a -> (MHGChains a, Log Double)
swapWith Int
i (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) MHGChains a
cs
  -- 2. Accept or reject.
  Bool
accept <- Log Double -> IOGenM StdGen -> IO Bool
mhgAccept Log Double
r IOGenM StdGen
g
  if Bool
accept
    then do
      -- -- Debug.
      -- traceIO $ "Swap accepted: " <> show i <> " <-> " <> show (i+1)
      -- let prL' = prior $ link $ fromMHG $ y V.! i
      --     prR' = prior $ link $ fromMHG $ y V.! (i+1)
      --     lhL' = likelihood $ link $ fromMHG $ y V.! i
      --     lhR' = likelihood $ link $ fromMHG $ y V.! (i+1)
      -- traceIO $ "Log priors (left, right, before swap): " <> show (ln prL) <> " " <> show (ln prR)
      -- traceIO $ "Log priors (left, right, after swap): " <> show (ln prL') <> " " <> show (ln prR')
      -- traceIO $ "Log likelihoods (left, right, before swap): " <> show (ln lhL) <> " " <> show (ln lhR)
      -- traceIO $ "Log likelihood (left, right, after swap): " <> show (ln lhL') <> " " <> show (ln lhR')
      let !ac' :: Acceptances Int
ac' = Maybe AcceptanceRates -> Int -> Acceptances Int -> Acceptances Int
forall k.
Ord k =>
Maybe AcceptanceRates -> k -> Acceptances k -> Acceptances k
pushAccept Maybe AcceptanceRates
forall a. Maybe a
Nothing Int
i (MC3 a -> Acceptances Int
forall a. MC3 a -> Acceptances Int
mc3SwapAcceptances MC3 a
a)
      MC3 a -> IO (MC3 a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (MC3 a -> IO (MC3 a)) -> MC3 a -> IO (MC3 a)
forall a b. (a -> b) -> a -> b
$ MC3 a
a {mc3MHGChains = y, mc3SwapAcceptances = ac'}
    else do
      let !ac' :: Acceptances Int
ac' = Maybe AcceptanceRates -> Int -> Acceptances Int -> Acceptances Int
forall k.
Ord k =>
Maybe AcceptanceRates -> k -> Acceptances k -> Acceptances k
pushReject Maybe AcceptanceRates
forall a. Maybe a
Nothing Int
i (MC3 a -> Acceptances Int
forall a. MC3 a -> Acceptances Int
mc3SwapAcceptances MC3 a
a)
      MC3 a -> IO (MC3 a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (MC3 a -> IO (MC3 a)) -> MC3 a -> IO (MC3 a)
forall a b. (a -> b) -> a -> b
$ MC3 a
a {mc3SwapAcceptances = ac'}
  where
    g :: IOGenM StdGen
g = MC3 a -> IOGenM StdGen
forall a. MC3 a -> IOGenM StdGen
mc3Generator MC3 a
a

mc3IsInvalidState :: (ToJSON a) => MC3 a -> Bool
mc3IsInvalidState :: forall a. ToJSON a => MC3 a -> Bool
mc3IsInvalidState MC3 a
a = (MHG a -> Bool) -> Vector (MHG a) -> Bool
forall a. (a -> Bool) -> Vector a -> Bool
V.any MHG a -> Bool
forall a. Algorithm a => a -> Bool
aIsInvalidState Vector (MHG a)
mhgs
  where
    mhgs :: Vector (MHG a)
mhgs = MC3 a -> Vector (MHG a)
forall a. MC3 a -> MHGChains a
mc3MHGChains MC3 a
a

-- NOTE: 'mc3Iterate' is actually not parallel, but concurrent because of the IO
-- constraint of the mutable trace.
mc3Iterate ::
  (ToJSON a) =>
  IterationMode ->
  ParallelizationMode ->
  MC3 a ->
  IO (MC3 a)
mc3Iterate :: forall a.
ToJSON a =>
IterationMode -> ParallelizationMode -> MC3 a -> IO (MC3 a)
mc3Iterate IterationMode
m ParallelizationMode
pm MC3 a
a = do
  -- 1. Maybe propose swaps.
  --
  -- NOTE: Swaps have to be proposed first, because the traces are automatically
  -- updated at step 2.
  let s :: MC3Settings
s = MC3 a -> MC3Settings
forall a. MC3 a -> MC3Settings
mc3Settings MC3 a
a
  MC3 a
a' <-
    if MC3 a -> Int
forall a. MC3 a -> Int
mc3Iteration MC3 a
a Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` SwapPeriod -> Int
fromSwapPeriod (MC3Settings -> SwapPeriod
mc3SwapPeriod MC3Settings
s) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
      then do
        let n :: Int
n = Vector (MHG a) -> Int
forall a. Vector a -> Int
V.length (Vector (MHG a) -> Int) -> Vector (MHG a) -> Int
forall a b. (a -> b) -> a -> b
$ MC3 a -> Vector (MHG a)
forall a. MC3 a -> MHGChains a
mc3MHGChains MC3 a
a
            is :: [Int]
is = [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2]
            ns :: Int
ns = NSwaps -> Int
fromNSwaps (NSwaps -> Int) -> NSwaps -> Int
forall a b. (a -> b) -> a -> b
$ MC3Settings -> NSwaps
mc3NSwaps MC3Settings
s
        [Int]
is' <- [Int] -> IOGenM StdGen -> IO [Int]
forall g (m :: * -> *) a. StatefulGen g m => [a] -> g -> m [a]
shuffle [Int]
is (IOGenM StdGen -> IO [Int]) -> IOGenM StdGen -> IO [Int]
forall a b. (a -> b) -> a -> b
$ MC3 a -> IOGenM StdGen
forall a. MC3 a -> IOGenM StdGen
mc3Generator MC3 a
a
        (MC3 a -> Int -> IO (MC3 a)) -> MC3 a -> [Int] -> IO (MC3 a)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM MC3 a -> Int -> IO (MC3 a)
forall a. MC3 a -> Int -> IO (MC3 a)
mc3ProposeSwap MC3 a
a (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
ns [Int]
is')
      else MC3 a -> IO (MC3 a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return MC3 a
a
  -- 2. Iterate all chains and increment iteration.
  Vector (MHG a)
mhgs <- case ParallelizationMode
pm of
    ParallelizationMode
Sequential -> (MHG a -> IO (MHG a)) -> Vector (MHG a) -> IO (Vector (MHG a))
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Vector a -> m (Vector b)
V.mapM (IterationMode -> ParallelizationMode -> MHG a -> IO (MHG a)
forall a.
Algorithm a =>
IterationMode -> ParallelizationMode -> a -> IO a
aIterate IterationMode
m ParallelizationMode
pm) (MC3 a -> Vector (MHG a)
forall a. MC3 a -> MHGChains a
mc3MHGChains MC3 a
a')
    ParallelizationMode
Parallel ->
      -- Go via a list, and use 'forkIO' ("Control.Concurrent.Async").
      [MHG a] -> Vector (MHG a)
forall a. [a] -> Vector a
V.fromList ([MHG a] -> Vector (MHG a)) -> IO [MHG a] -> IO (Vector (MHG a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (MHG a -> IO (MHG a)) -> [MHG a] -> IO [MHG a]
forall (t :: * -> *) a b.
Traversable t =>
(a -> IO b) -> t a -> IO (t b)
mapConcurrently (IterationMode -> ParallelizationMode -> MHG a -> IO (MHG a)
forall a.
Algorithm a =>
IterationMode -> ParallelizationMode -> a -> IO a
aIterate IterationMode
m ParallelizationMode
pm) (Vector (MHG a) -> [MHG a]
forall a. Vector a -> [a]
V.toList (Vector (MHG a) -> [MHG a]) -> Vector (MHG a) -> [MHG a]
forall a b. (a -> b) -> a -> b
$ MC3 a -> Vector (MHG a)
forall a. MC3 a -> MHGChains a
mc3MHGChains MC3 a
a')
  let i :: Int
i = MC3 a -> Int
forall a. MC3 a -> Int
mc3Iteration MC3 a
a'
  MC3 a -> IO (MC3 a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (MC3 a -> IO (MC3 a)) -> MC3 a -> IO (MC3 a)
forall a b. (a -> b) -> a -> b
$ MC3 a
a' {mc3MHGChains = mhgs, mc3Iteration = succ i}

tuneBeta ::
  -- The old reciprocal temperatures are needed to retrieve the old ratios.
  ReciprocalTemperatures ->
  -- Index i of left chain. Change the reciprocal temperature of chain (i+1).
  Int ->
  -- Exponent xi of the reciprocal temperature ratio.
  Double ->
  -- The new reciprocal temperatures are updated incrementally using the
  -- reciprocal temperature ratios during the fold (see 'mc3AutoTune' below).
  ReciprocalTemperatures ->
  ReciprocalTemperatures
tuneBeta :: ReciprocalTemperatures
-> Int
-> Double
-> ReciprocalTemperatures
-> ReciprocalTemperatures
tuneBeta ReciprocalTemperatures
bsOld Int
i Double
xi ReciprocalTemperatures
bsNew = ReciprocalTemperatures
bsNew ReciprocalTemperatures -> [(Int, Double)] -> ReciprocalTemperatures
forall a. Unbox a => Vector a -> [(Int, a)] -> Vector a
U.// [(Int
j, Double
brNew)]
  where
    j :: Int
j = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
    blOld :: Double
blOld = ReciprocalTemperatures
bsOld ReciprocalTemperatures -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
U.! Int
i
    brOld :: Double
brOld = ReciprocalTemperatures
bsOld ReciprocalTemperatures -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
U.! Int
j
    blNew :: Double
blNew = ReciprocalTemperatures
bsNew ReciprocalTemperatures -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
U.! Int
i
    -- The new ratio is in (0,1).
    rNew :: Double
rNew = (Double
brOld Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
blOld) Double -> Double -> Double
forall a. Floating a => a -> a -> a
** Double
xi
    brNew :: Double
brNew = Double
blNew Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
rNew

mc3AutoTune :: (ToJSON a) => TuningType -> Int -> MC3 a -> IO (MC3 a)
mc3AutoTune :: forall a. ToJSON a => TuningType -> Int -> MC3 a -> IO (MC3 a)
mc3AutoTune TuningType
b Int
l MC3 a
a = do
  -- 1. Auto tune all chains.
  Vector (MHG a)
mhgs' <- (MHG a -> IO (MHG a)) -> Vector (MHG a) -> IO (Vector (MHG a))
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Vector a -> m (Vector b)
V.mapM (TuningType -> Int -> MHG a -> IO (MHG a)
forall a. Algorithm a => TuningType -> Int -> a -> IO a
aAutoTune TuningType
b Int
l) (Vector (MHG a) -> IO (Vector (MHG a)))
-> Vector (MHG a) -> IO (Vector (MHG a))
forall a b. (a -> b) -> a -> b
$ MC3 a -> Vector (MHG a)
forall a. MC3 a -> MHGChains a
mc3MHGChains MC3 a
a
  -- 2. Auto tune temperatures.
  let optimalRate :: Double
optimalRate = PDimension -> Double
getOptimalRate PDimension
PDimensionUnknown
      mCurrentRates :: Map Int (Maybe Double)
mCurrentRates = Acceptances Int -> Map Int (Maybe Double)
forall k. Acceptances k -> Map k (Maybe Double)
acceptanceRates (Acceptances Int -> Map Int (Maybe Double))
-> Acceptances Int -> Map Int (Maybe Double)
forall a b. (a -> b) -> a -> b
$ MC3 a -> Acceptances Int
forall a. MC3 a -> Acceptances Int
mc3SwapAcceptances MC3 a
a
      -- We assume that the acceptance rate of state swaps between two chains is
      -- roughly proportional to the ratio of the temperatures of the chains.
      -- Hence, we focus on temperature ratios, actually reciprocal temperature
      -- ratios, which is the same. Also, by working with ratios in (0,1) of
      -- neighboring chains, we ensure the monotonicity of the reciprocal
      -- temperatures.
      --
      -- The factor (1/2) was determined by a few tests and is otherwise
      -- absolutely arbitrary.
      xi :: Int -> Double
xi Int
i = case Map Int (Maybe Double)
mCurrentRates Map Int (Maybe Double) -> Int -> Maybe Double
forall k a. Ord k => Map k a -> k -> a
M.! Int
i of
        Maybe Double
Nothing -> Double
1.0
        Just Double
currentRate -> Double -> Double
forall a. Floating a => a -> a
exp (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ (Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
2) (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ Double
currentRate Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
optimalRate
      bs :: ReciprocalTemperatures
bs = MC3 a -> ReciprocalTemperatures
forall a. MC3 a -> ReciprocalTemperatures
mc3ReciprocalTemperatures MC3 a
a
      n :: Int
n = NChains -> Int
fromNChains (NChains -> Int) -> NChains -> Int
forall a b. (a -> b) -> a -> b
$ MC3Settings -> NChains
mc3NChains (MC3Settings -> NChains) -> MC3Settings -> NChains
forall a b. (a -> b) -> a -> b
$ MC3 a -> MC3Settings
forall a. MC3 a -> MC3Settings
mc3Settings MC3 a
a
      -- Do not change the temperature, and the prior and likelihood functions of
      -- the cold chain.
      bs' :: ReciprocalTemperatures
bs' = (ReciprocalTemperatures -> Int -> ReciprocalTemperatures)
-> ReciprocalTemperatures -> [Int] -> ReciprocalTemperatures
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\ReciprocalTemperatures
xs Int
j -> ReciprocalTemperatures
-> Int
-> Double
-> ReciprocalTemperatures
-> ReciprocalTemperatures
tuneBeta ReciprocalTemperatures
bs Int
j (Int -> Double
xi Int
j) ReciprocalTemperatures
xs) ReciprocalTemperatures
bs [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2]
      coldChain :: Chain a
coldChain = MHG a -> Chain a
forall a. MHG a -> Chain a
fromMHG (MHG a -> Chain a) -> MHG a -> Chain a
forall a b. (a -> b) -> a -> b
$ Vector (MHG a) -> MHG a
forall a. Vector a -> a
V.head Vector (MHG a)
mhgs'
      coldPrF :: PriorFunction a
coldPrF = Chain a -> PriorFunction a
forall a. Chain a -> PriorFunction a
priorFunction Chain a
coldChain
      coldLhF :: PriorFunction a
coldLhF = Chain a -> PriorFunction a
forall a. Chain a -> PriorFunction a
likelihoodFunction Chain a
coldChain
      mhgs'' :: Vector (MHG a)
mhgs'' =
        Vector (MHG a) -> MHG a
forall a. Vector a -> a
V.head Vector (MHG a)
mhgs'
          MHG a -> Vector (MHG a) -> Vector (MHG a)
forall a. a -> Vector a -> Vector a
`V.cons` (Double -> MHG a -> MHG a)
-> Vector Double -> Vector (MHG a) -> Vector (MHG a)
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
V.zipWith
            (PriorFunction a -> PriorFunction a -> Double -> MHG a -> MHG a
forall a.
PriorFunction a -> PriorFunction a -> Double -> MHG a -> MHG a
setReciprocalTemperature PriorFunction a
coldPrF PriorFunction a
coldLhF)
            (ReciprocalTemperatures -> Vector Double
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
V.convert (ReciprocalTemperatures -> Vector Double)
-> ReciprocalTemperatures -> Vector Double
forall a b. (a -> b) -> a -> b
$ ReciprocalTemperatures -> ReciprocalTemperatures
forall a. Unbox a => Vector a -> Vector a
U.tail ReciprocalTemperatures
bs')
            (Vector (MHG a) -> Vector (MHG a)
forall a. Vector a -> Vector a
V.tail Vector (MHG a)
mhgs')
  MC3 a -> IO (MC3 a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (MC3 a -> IO (MC3 a)) -> MC3 a -> IO (MC3 a)
forall a b. (a -> b) -> a -> b
$
    if TuningType
b TuningType -> TuningType -> Bool
forall a. Eq a => a -> a -> Bool
== TuningType
NormalTuningFastProposalsOnly Bool -> Bool -> Bool
|| TuningType
b TuningType -> TuningType -> Bool
forall a. Eq a => a -> a -> Bool
== TuningType
NormalTuningAllProposals
      then MC3 a
a {mc3MHGChains = mhgs'', mc3ReciprocalTemperatures = bs'}
      else MC3 a
a {mc3MHGChains = mhgs'}

mc3ResetAcceptance :: (ToJSON a) => ResetAcceptance -> MC3 a -> MC3 a
mc3ResetAcceptance :: forall a. ToJSON a => ResetAcceptance -> MC3 a -> MC3 a
mc3ResetAcceptance ResetAcceptance
x MC3 a
a = MC3 a
a'
  where
    -- 1. Reset acceptance of all chains.
    mhgs' :: Vector (MHG a)
mhgs' = (MHG a -> MHG a) -> Vector (MHG a) -> Vector (MHG a)
forall a b. (a -> b) -> Vector a -> Vector b
V.map (ResetAcceptance -> MHG a -> MHG a
forall a. Algorithm a => ResetAcceptance -> a -> a
aResetAcceptance ResetAcceptance
x) (MC3 a -> Vector (MHG a)
forall a. MC3 a -> MHGChains a
mc3MHGChains MC3 a
a)
    -- 2. Reset acceptance of swaps.
    ac' :: Acceptances Int
ac' = ResetAcceptance -> Acceptances Int -> Acceptances Int
forall k.
Ord k =>
ResetAcceptance -> Acceptances k -> Acceptances k
resetA ResetAcceptance
x (Acceptances Int -> Acceptances Int)
-> Acceptances Int -> Acceptances Int
forall a b. (a -> b) -> a -> b
$ MC3 a -> Acceptances Int
forall a. MC3 a -> Acceptances Int
mc3SwapAcceptances MC3 a
a
    --
    a' :: MC3 a
a' = MC3 a
a {mc3MHGChains = mhgs', mc3SwapAcceptances = ac'}

mc3CleanAfterBurnIn :: (ToJSON a) => TraceLength -> MC3 a -> IO (MC3 a)
mc3CleanAfterBurnIn :: forall a. ToJSON a => TraceLength -> MC3 a -> IO (MC3 a)
mc3CleanAfterBurnIn TraceLength
tl MC3 a
a = do
  Vector (MHG a)
cs' <- (MHG a -> IO (MHG a)) -> Vector (MHG a) -> IO (Vector (MHG a))
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Vector a -> m (Vector b)
V.mapM (TraceLength -> MHG a -> IO (MHG a)
forall a. Algorithm a => TraceLength -> a -> IO a
aCleanAfterBurnIn TraceLength
tl) Vector (MHG a)
cs
  MC3 a -> IO (MC3 a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MC3 a -> IO (MC3 a)) -> MC3 a -> IO (MC3 a)
forall a b. (a -> b) -> a -> b
$ MC3 a
a {mc3MHGChains = cs'}
  where
    cs :: Vector (MHG a)
cs = MC3 a -> Vector (MHG a)
forall a. MC3 a -> MHGChains a
mc3MHGChains MC3 a
a

-- Information in cycle summary:
--
-- - The complete summary of the cycle of the cold chain.
--
-- - The combined acceptance rate of proposals within the hot chains.
--
-- - The temperatures of the chains and the acceptance rates of the state swaps.
mc3SummarizeCycle :: (ToJSON a) => IterationMode -> MC3 a -> BL.ByteString
mc3SummarizeCycle :: forall a. ToJSON a => IterationMode -> MC3 a -> ByteString
mc3SummarizeCycle IterationMode
m MC3 a
a =
  ByteString -> [ByteString] -> ByteString
BL.intercalate ByteString
"\n" ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$
    [ ByteString
"MC3: Cycle of cold chain.",
      ByteString
coldMHGCycleSummary
    ]
      [ByteString] -> [ByteString] -> [ByteString]
forall a. [a] -> [a] -> [a]
++ case Maybe Double
mAr of
        Maybe Double
Nothing -> []
        Just Double
ar ->
          [ ByteString
"MC3: Average acceptance rate across all chains: "
              ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> Builder -> ByteString
BB.toLazyByteString (FloatFormat -> Double -> Builder
BB.formatDouble (Int -> FloatFormat
BB.standard Int
2) Double
ar)
              ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
"."
          ]
      [ByteString] -> [ByteString] -> [ByteString]
forall a. [a] -> [a] -> [a]
++ [ ByteString
"MC3: Reciprocal temperatures of the chains: " ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString -> [ByteString] -> ByteString
BL.intercalate ByteString
", " [ByteString]
bsB ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
".",
           ByteString
"MC3: Summary of state swaps.",
           ByteString
"MC3: The swap period is " ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
swapPeriodB ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
".",
           ByteString
"MC3: The state swaps are executed in random order.",
           ByteString
proposalHeader,
           ByteString
proposalHLine
         ]
      [ByteString] -> [ByteString] -> [ByteString]
forall a. [a] -> [a] -> [a]
++ [ PName
-> PDescription
-> PWeight
-> Maybe Double
-> PDimension
-> (Int, Int, Maybe Double, Maybe Double)
-> ByteString
summarizeProposal
             ([Char] -> PName
PName ([Char] -> PName) -> [Char] -> PName
forall a b. (a -> b) -> a -> b
$ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
i [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ [Char]
" <-> " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
             ([Char] -> PDescription
PDescription [Char]
"Swap states between chains")
             (Int -> PWeight
pWeight Int
1)
             (Double -> Maybe Double
forall a. a -> Maybe a
Just (Double -> Maybe Double) -> Double -> Maybe Double
forall a b. (a -> b) -> a -> b
$ ReciprocalTemperatures
bs ReciprocalTemperatures -> Int -> Double
forall a. Unbox a => Vector a -> Int -> a
U.! (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
             PDimension
PDimensionUnknown
             (Int -> Acceptances Int -> (Int, Int, Maybe Double, Maybe Double)
forall k.
Ord k =>
k -> Acceptances k -> (Int, Int, Maybe Double, Maybe Double)
acceptanceRate Int
i Acceptances Int
swapAcceptance)
           | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2]
         ]
      [ByteString] -> [ByteString] -> [ByteString]
forall a. [a] -> [a] -> [a]
++ [ByteString
proposalHLine]
  where
    mhgs :: MHGChains a
mhgs = MC3 a -> MHGChains a
forall a. MC3 a -> MHGChains a
mc3MHGChains MC3 a
a
    coldMHGCycleSummary :: ByteString
coldMHGCycleSummary = IterationMode -> MHG a -> ByteString
forall a. Algorithm a => IterationMode -> a -> ByteString
aSummarizeCycle IterationMode
m (MHG a -> ByteString) -> MHG a -> ByteString
forall a b. (a -> b) -> a -> b
$ MHGChains a -> MHG a
forall a. Vector a -> a
V.head MHGChains a
mhgs
    cs :: Vector (Chain a)
cs = (MHG a -> Chain a) -> MHGChains a -> Vector (Chain a)
forall a b. (a -> b) -> Vector a -> Vector b
V.map MHG a -> Chain a
forall a. MHG a -> Chain a
fromMHG MHGChains a
mhgs
    -- Acceptance rates may be 'Nothing' when no proposals have been undertaken.
    -- The 'sequence' operations pull the 'Nothing's out of the inner
    -- structures.
    as :: Maybe (Vector (Map (Proposal a) Double))
as = Vector (Maybe (Map (Proposal a) Double))
-> Maybe (Vector (Map (Proposal a) Double))
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => Vector (m a) -> m (Vector a)
sequence (Vector (Maybe (Map (Proposal a) Double))
 -> Maybe (Vector (Map (Proposal a) Double)))
-> Vector (Maybe (Map (Proposal a) Double))
-> Maybe (Vector (Map (Proposal a) Double))
forall a b. (a -> b) -> a -> b
$ (Chain a -> Maybe (Map (Proposal a) Double))
-> Vector (Chain a) -> Vector (Maybe (Map (Proposal a) Double))
forall a b. (a -> b) -> Vector a -> Vector b
V.map (Map (Proposal a) (Maybe Double) -> Maybe (Map (Proposal a) Double)
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a.
Monad m =>
Map (Proposal a) (m a) -> m (Map (Proposal a) a)
sequence (Map (Proposal a) (Maybe Double)
 -> Maybe (Map (Proposal a) Double))
-> (Chain a -> Map (Proposal a) (Maybe Double))
-> Chain a
-> Maybe (Map (Proposal a) Double)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Acceptances (Proposal a) -> Map (Proposal a) (Maybe Double)
forall k. Acceptances k -> Map k (Maybe Double)
acceptanceRates (Acceptances (Proposal a) -> Map (Proposal a) (Maybe Double))
-> (Chain a -> Acceptances (Proposal a))
-> Chain a
-> Map (Proposal a) (Maybe Double)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Chain a -> Acceptances (Proposal a)
forall a. Chain a -> Acceptances (Proposal a)
acceptances) Vector (Chain a)
cs
    mVecAr :: Maybe (Vector Double)
mVecAr = (Map (Proposal a) Double -> Double)
-> Vector (Map (Proposal a) Double) -> Vector Double
forall a b. (a -> b) -> Vector a -> Vector b
V.map (\Map (Proposal a) Double
mp -> Map (Proposal a) Double -> Double
forall a. Num a => Map (Proposal a) a -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum Map (Proposal a) Double
mp Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Map (Proposal a) Double -> Int
forall a. Map (Proposal a) a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Map (Proposal a) Double
mp)) (Vector (Map (Proposal a) Double) -> Vector Double)
-> Maybe (Vector (Map (Proposal a) Double))
-> Maybe (Vector Double)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (Vector (Map (Proposal a) Double))
as
    mAr :: Maybe Double
mAr = (\Vector Double
vec -> Vector Double -> Double
forall a. Num a => Vector a -> a
V.sum Vector Double
vec Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Vector Double -> Int
forall a. Vector a -> Int
V.length Vector Double
vec)) (Vector Double -> Double) -> Maybe (Vector Double) -> Maybe Double
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (Vector Double)
mVecAr
    bs :: ReciprocalTemperatures
bs = MC3 a -> ReciprocalTemperatures
forall a. MC3 a -> ReciprocalTemperatures
mc3ReciprocalTemperatures MC3 a
a
    bsB :: [ByteString]
bsB = (Double -> ByteString) -> [Double] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
map (Builder -> ByteString
BB.toLazyByteString (Builder -> ByteString)
-> (Double -> Builder) -> Double -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FloatFormat -> Double -> Builder
BB.formatDouble (Int -> FloatFormat
BB.standard Int
2)) ([Double] -> [ByteString]) -> [Double] -> [ByteString]
forall a b. (a -> b) -> a -> b
$ ReciprocalTemperatures -> [Double]
forall a. Unbox a => Vector a -> [a]
U.toList ReciprocalTemperatures
bs
    swapPeriod :: Int
swapPeriod = SwapPeriod -> Int
fromSwapPeriod (SwapPeriod -> Int) -> SwapPeriod -> Int
forall a b. (a -> b) -> a -> b
$ MC3Settings -> SwapPeriod
mc3SwapPeriod (MC3Settings -> SwapPeriod) -> MC3Settings -> SwapPeriod
forall a b. (a -> b) -> a -> b
$ MC3 a -> MC3Settings
forall a. MC3 a -> MC3Settings
mc3Settings MC3 a
a
    swapPeriodB :: ByteString
swapPeriodB = Builder -> ByteString
BB.toLazyByteString (Builder -> ByteString) -> Builder -> ByteString
forall a b. (a -> b) -> a -> b
$ Int -> Builder
BB.intDec Int
swapPeriod
    swapAcceptance :: Acceptances Int
swapAcceptance = MC3 a -> Acceptances Int
forall a. MC3 a -> Acceptances Int
mc3SwapAcceptances MC3 a
a
    n :: Int
n = NChains -> Int
fromNChains (NChains -> Int) -> NChains -> Int
forall a b. (a -> b) -> a -> b
$ MC3Settings -> NChains
mc3NChains (MC3Settings -> NChains) -> MC3Settings -> NChains
forall a b. (a -> b) -> a -> b
$ MC3 a -> MC3Settings
forall a. MC3 a -> MC3Settings
mc3Settings MC3 a
a
    proposalHLine :: ByteString
proposalHLine = Int64 -> Char -> ByteString
BL.replicate (ByteString -> Int64
BL.length ByteString
proposalHeader) Char
'-'

-- No extra monitors are opened.
mc3OpenMonitors :: (ToJSON a) => AnalysisName -> ExecutionMode -> MC3 a -> IO (MC3 a)
mc3OpenMonitors :: forall a.
ToJSON a =>
AnalysisName -> ExecutionMode -> MC3 a -> IO (MC3 a)
mc3OpenMonitors AnalysisName
nm ExecutionMode
em MC3 a
a = do
  Vector (MHG a)
mhgs' <- (Int -> MHG a -> IO (MHG a))
-> Vector (MHG a) -> IO (Vector (MHG a))
forall (m :: * -> *) a b.
Monad m =>
(Int -> a -> m b) -> Vector a -> m (Vector b)
V.imapM Int -> MHG a -> IO (MHG a)
forall {p} {a}. PrintfArg p => p -> MHG a -> IO (MHG a)
mhgOpenMonitors (MC3 a -> Vector (MHG a)
forall a. MC3 a -> MHGChains a
mc3MHGChains MC3 a
a)
  MC3 a -> IO (MC3 a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (MC3 a -> IO (MC3 a)) -> MC3 a -> IO (MC3 a)
forall a b. (a -> b) -> a -> b
$ MC3 a
a {mc3MHGChains = mhgs'}
  where
    mhgOpenMonitors :: p -> MHG a -> IO (MHG a)
mhgOpenMonitors p
i (MHG Chain a
c) = do
      Monitor a
m' <- [Char] -> [Char] -> ExecutionMode -> Monitor a -> IO (Monitor a)
forall a.
[Char] -> [Char] -> ExecutionMode -> Monitor a -> IO (Monitor a)
mOpen [Char]
pre [Char]
suf ExecutionMode
em (Monitor a -> IO (Monitor a)) -> Monitor a -> IO (Monitor a)
forall a b. (a -> b) -> a -> b
$ Chain a -> Monitor a
forall a. Chain a -> Monitor a
monitor Chain a
c
      MHG a -> IO (MHG a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MHG a -> IO (MHG a)) -> MHG a -> IO (MHG a)
forall a b. (a -> b) -> a -> b
$ Chain a -> MHG a
forall a. Chain a -> MHG a
MHG Chain a
c {monitor = m'}
      where
        pre :: [Char]
pre = AnalysisName -> [Char]
fromAnalysisName AnalysisName
nm
        suf :: [Char]
suf = [Char] -> p -> [Char]
forall r. PrintfType r => [Char] -> r
printf [Char]
"%02d" p
i

mc3ExecuteMonitors ::
  (ToJSON a) =>
  Verbosity ->
  -- Starting time.
  UTCTime ->
  -- Total number of iterations.
  Int ->
  MC3 a ->
  IO (Maybe BL.ByteString)
mc3ExecuteMonitors :: forall a.
ToJSON a =>
Verbosity -> UTCTime -> Int -> MC3 a -> IO (Maybe ByteString)
mc3ExecuteMonitors Verbosity
vb UTCTime
t0 Int
iTotal MC3 a
a = Vector (Maybe ByteString) -> Maybe ByteString
forall a. Vector a -> a
V.head (Vector (Maybe ByteString) -> Maybe ByteString)
-> IO (Vector (Maybe ByteString)) -> IO (Maybe ByteString)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int -> MHG a -> IO (Maybe ByteString))
-> Vector (MHG a) -> IO (Vector (Maybe ByteString))
forall (m :: * -> *) a b.
Monad m =>
(Int -> a -> m b) -> Vector a -> m (Vector b)
V.imapM Int -> MHG a -> IO (Maybe ByteString)
forall {a} {a}.
(Eq a, Num a, Algorithm a) =>
a -> a -> IO (Maybe ByteString)
f (MC3 a -> Vector (MHG a)
forall a. MC3 a -> MHGChains a
mc3MHGChains MC3 a
a)
  where
    -- The first chain honors verbosity.
    f :: a -> a -> IO (Maybe ByteString)
f a
0 = Verbosity -> UTCTime -> Int -> a -> IO (Maybe ByteString)
forall a.
Algorithm a =>
Verbosity -> UTCTime -> Int -> a -> IO (Maybe ByteString)
aExecuteMonitors Verbosity
vb UTCTime
t0 Int
iTotal
    -- All other chains are to be quiet.
    f a
_ = Verbosity -> UTCTime -> Int -> a -> IO (Maybe ByteString)
forall a.
Algorithm a =>
Verbosity -> UTCTime -> Int -> a -> IO (Maybe ByteString)
aExecuteMonitors Verbosity
Quiet UTCTime
t0 Int
iTotal

mc3StdMonitorHeader :: (ToJSON a) => MC3 a -> BL.ByteString
mc3StdMonitorHeader :: forall a. ToJSON a => MC3 a -> ByteString
mc3StdMonitorHeader = MHG a -> ByteString
forall a. Algorithm a => a -> ByteString
aStdMonitorHeader (MHG a -> ByteString) -> (MC3 a -> MHG a) -> MC3 a -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (MHG a) -> MHG a
forall a. Vector a -> a
V.head (Vector (MHG a) -> MHG a)
-> (MC3 a -> Vector (MHG a)) -> MC3 a -> MHG a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MC3 a -> Vector (MHG a)
forall a. MC3 a -> MHGChains a
mc3MHGChains

mc3CloseMonitors :: (ToJSON a) => MC3 a -> IO (MC3 a)
mc3CloseMonitors :: forall a. ToJSON a => MC3 a -> IO (MC3 a)
mc3CloseMonitors MC3 a
a = do
  Vector (MHG a)
mhgs' <- (MHG a -> IO (MHG a)) -> Vector (MHG a) -> IO (Vector (MHG a))
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Vector a -> m (Vector b)
V.mapM MHG a -> IO (MHG a)
forall a. Algorithm a => a -> IO a
aCloseMonitors (Vector (MHG a) -> IO (Vector (MHG a)))
-> Vector (MHG a) -> IO (Vector (MHG a))
forall a b. (a -> b) -> a -> b
$ MC3 a -> Vector (MHG a)
forall a. MC3 a -> MHGChains a
mc3MHGChains MC3 a
a
  MC3 a -> IO (MC3 a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (MC3 a -> IO (MC3 a)) -> MC3 a -> IO (MC3 a)
forall a b. (a -> b) -> a -> b
$ MC3 a
a {mc3MHGChains = mhgs'}