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

-- |
-- Module      :  Mcmc.Algorithm.MC3
-- Description :  Metropolis-coupled Markov chain Monte Carlo algorithm
-- Copyright   :  (c) Dominik Schrempf, 2021
-- 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 qualified Data.Double.Conversion.ByteString as BC
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.MWC

-- import Debug.Trace hiding (trace)

-- | 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
/= :: NChains -> NChains -> Bool
$c/= :: NChains -> NChains -> Bool
== :: NChains -> NChains -> Bool
$c== :: 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
readListPrec :: ReadPrec [NChains]
$creadListPrec :: ReadPrec [NChains]
readPrec :: ReadPrec NChains
$creadPrec :: ReadPrec NChains
readList :: ReadS [NChains]
$creadList :: ReadS [NChains]
readsPrec :: Int -> ReadS NChains
$creadsPrec :: Int -> ReadS NChains
Read, Int -> NChains -> ShowS
[NChains] -> ShowS
NChains -> String
(Int -> NChains -> ShowS)
-> (NChains -> String) -> ([NChains] -> ShowS) -> Show NChains
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [NChains] -> ShowS
$cshowList :: [NChains] -> ShowS
show :: NChains -> String
$cshow :: NChains -> String
showsPrec :: Int -> NChains -> ShowS
$cshowsPrec :: Int -> 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
/= :: SwapPeriod -> SwapPeriod -> Bool
$c/= :: SwapPeriod -> SwapPeriod -> Bool
== :: SwapPeriod -> SwapPeriod -> Bool
$c== :: 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
readListPrec :: ReadPrec [SwapPeriod]
$creadListPrec :: ReadPrec [SwapPeriod]
readPrec :: ReadPrec SwapPeriod
$creadPrec :: ReadPrec SwapPeriod
readList :: ReadS [SwapPeriod]
$creadList :: ReadS [SwapPeriod]
readsPrec :: Int -> ReadS SwapPeriod
$creadsPrec :: Int -> ReadS SwapPeriod
Read, Int -> SwapPeriod -> ShowS
[SwapPeriod] -> ShowS
SwapPeriod -> String
(Int -> SwapPeriod -> ShowS)
-> (SwapPeriod -> String)
-> ([SwapPeriod] -> ShowS)
-> Show SwapPeriod
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [SwapPeriod] -> ShowS
$cshowList :: [SwapPeriod] -> ShowS
show :: SwapPeriod -> String
$cshow :: SwapPeriod -> String
showsPrec :: Int -> SwapPeriod -> ShowS
$cshowsPrec :: Int -> 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
/= :: NSwaps -> NSwaps -> Bool
$c/= :: NSwaps -> NSwaps -> Bool
== :: NSwaps -> NSwaps -> Bool
$c== :: 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
readListPrec :: ReadPrec [NSwaps]
$creadListPrec :: ReadPrec [NSwaps]
readPrec :: ReadPrec NSwaps
$creadPrec :: ReadPrec NSwaps
readList :: ReadS [NSwaps]
$creadList :: ReadS [NSwaps]
readsPrec :: Int -> ReadS NSwaps
$creadsPrec :: Int -> ReadS NSwaps
Read, Int -> NSwaps -> ShowS
[NSwaps] -> ShowS
NSwaps -> String
(Int -> NSwaps -> ShowS)
-> (NSwaps -> String) -> ([NSwaps] -> ShowS) -> Show NSwaps
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [NSwaps] -> ShowS
$cshowList :: [NSwaps] -> ShowS
show :: NSwaps -> String
$cshow :: NSwaps -> String
showsPrec :: Int -> NSwaps -> ShowS
$cshowsPrec :: Int -> 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
/= :: MC3Settings -> MC3Settings -> Bool
$c/= :: MC3Settings -> MC3Settings -> Bool
== :: MC3Settings -> MC3Settings -> Bool
$c== :: 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
readListPrec :: ReadPrec [MC3Settings]
$creadListPrec :: ReadPrec [MC3Settings]
readPrec :: ReadPrec MC3Settings
$creadPrec :: ReadPrec MC3Settings
readList :: ReadS [MC3Settings]
$creadList :: ReadS [MC3Settings]
readsPrec :: Int -> ReadS MC3Settings
$creadsPrec :: Int -> ReadS MC3Settings
Read, Int -> MC3Settings -> ShowS
[MC3Settings] -> ShowS
MC3Settings -> String
(Int -> MC3Settings -> ShowS)
-> (MC3Settings -> String)
-> ([MC3Settings] -> ShowS)
-> Show MC3Settings
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [MC3Settings] -> ShowS
$cshowList :: [MC3Settings] -> ShowS
show :: MC3Settings -> String
$cshow :: MC3Settings -> String
showsPrec :: Int -> MC3Settings -> ShowS
$cshowsPrec :: Int -> 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
  { SavedMC3 a -> MC3Settings
savedMC3Settings :: MC3Settings,
    SavedMC3 a -> Vector (SavedChain a)
savedMC3Chains :: V.Vector (SavedChain a),
    SavedMC3 a -> ReciprocalTemperatures
savedMC3ReciprocalTemperatures :: ReciprocalTemperatures,
    SavedMC3 a -> Int
savedMC3Iteration :: Int,
    SavedMC3 a -> Acceptance Int
savedMC3SwapAcceptance :: Acceptance Int,
    SavedMC3 a -> Vector Word32
savedMC3Generator :: U.Vector Word32
  }
  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
/= :: SavedMC3 a -> SavedMC3 a -> Bool
$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
Eq, ReadPrec [SavedMC3 a]
ReadPrec (SavedMC3 a)
Int -> ReadS (SavedMC3 a)
ReadS [SavedMC3 a]
(Int -> ReadS (SavedMC3 a))
-> ReadS [SavedMC3 a]
-> ReadPrec (SavedMC3 a)
-> ReadPrec [SavedMC3 a]
-> Read (SavedMC3 a)
forall a. Read a => ReadPrec [SavedMC3 a]
forall a. Read a => ReadPrec (SavedMC3 a)
forall a. Read a => Int -> ReadS (SavedMC3 a)
forall a. Read a => ReadS [SavedMC3 a]
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
readListPrec :: ReadPrec [SavedMC3 a]
$creadListPrec :: forall a. Read a => ReadPrec [SavedMC3 a]
readPrec :: ReadPrec (SavedMC3 a)
$creadPrec :: forall a. Read a => ReadPrec (SavedMC3 a)
readList :: ReadS [SavedMC3 a]
$creadList :: forall a. Read a => ReadS [SavedMC3 a]
readsPrec :: Int -> ReadS (SavedMC3 a)
$creadsPrec :: forall a. Read a => Int -> ReadS (SavedMC3 a)
Read, Int -> SavedMC3 a -> ShowS
[SavedMC3 a] -> ShowS
SavedMC3 a -> String
(Int -> SavedMC3 a -> ShowS)
-> (SavedMC3 a -> String)
-> ([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 -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [SavedMC3 a] -> ShowS
$cshowList :: forall a. Show a => [SavedMC3 a] -> ShowS
show :: SavedMC3 a -> String
$cshow :: forall a. Show a => SavedMC3 a -> String
showsPrec :: Int -> SavedMC3 a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> SavedMC3 a -> ShowS
Show)

$(deriveJSON defaultOptions ''SavedMC3)

toSavedMC3 ::
  MC3 a ->
  IO (SavedMC3 a)
toSavedMC3 :: MC3 a -> IO (SavedMC3 a)
toSavedMC3 (MC3 MC3Settings
s MHGChains a
mhgs ReciprocalTemperatures
bs Int
i Acceptance Int
ac GenIO
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
  Vector Word32
g' <- GenIO -> IO (Vector Word32)
saveGen GenIO
g
  SavedMC3 a -> IO (SavedMC3 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
-> Acceptance Int
-> Vector Word32
-> SavedMC3 a
forall a.
MC3Settings
-> Vector (SavedChain a)
-> ReciprocalTemperatures
-> Int
-> Acceptance Int
-> Vector Word32
-> SavedMC3 a
SavedMC3 MC3Settings
s Vector (SavedChain a)
scs ReciprocalTemperatures
bs Int
i Acceptance Int
ac Vector Word32
g'

fromSavedMC3 ::
  PriorFunction a ->
  LikelihoodFunction a ->
  Cycle a ->
  Monitor a ->
  SavedMC3 a ->
  IO (MC3 a)
fromSavedMC3 :: 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 Acceptance Int
ac Vector Word32
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)
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
        ]
  Gen RealWorld
g <- Vector Word32 -> IO GenIO
loadGen Vector Word32
g'
  MC3 a -> IO (MC3 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
-> Acceptance Int
-> GenIO
-> MC3 a
forall a.
MC3Settings
-> MHGChains a
-> ReciprocalTemperatures
-> Int
-> Acceptance Int
-> GenIO
-> MC3 a
MC3 MC3Settings
s Vector (MHG a)
mhgs ReciprocalTemperatures
bs Int
i Acceptance Int
ac Gen RealWorld
GenIO
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
  { MC3 a -> MC3Settings
mc3Settings :: MC3Settings,
    -- | The first chain is the cold chain with temperature 1.0.
    MC3 a -> MHGChains a
mc3MHGChains :: MHGChains a,
    -- | Vector of reciprocal temperatures.
    MC3 a -> ReciprocalTemperatures
mc3ReciprocalTemperatures :: ReciprocalTemperatures,
    -- | Current iteration.
    MC3 a -> Int
mc3Iteration :: Int,
    -- | Number of accepted and rejected swaps.
    MC3 a -> Acceptance Int
mc3SwapAcceptance :: Acceptance Int,
    MC3 a -> GenIO
mc3Generator :: GenIO
  }

instance ToJSON a => Algorithm (MC3 a) where
  aName :: MC3 a -> String
aName = String -> MC3 a -> String
forall a b. a -> b -> a
const String
"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 :: ParallelizationMode -> MC3 a -> IO (MC3 a)
aIterate = ParallelizationMode -> MC3 a -> IO (MC3 a)
forall a. ToJSON a => ParallelizationMode -> MC3 a -> IO (MC3 a)
mc3Iterate
  aAutoTune :: Int -> MC3 a -> IO (MC3 a)
aAutoTune = Int -> MC3 a -> IO (MC3 a)
forall a. ToJSON a => Int -> MC3 a -> IO (MC3 a)
mc3AutoTune
  aResetAcceptance :: MC3 a -> MC3 a
aResetAcceptance = MC3 a -> MC3 a
forall a. ToJSON a => MC3 a -> MC3 a
mc3ResetAcceptance
  aSummarizeCycle :: MC3 a -> ByteString
aSummarizeCycle = MC3 a -> ByteString
forall a. ToJSON a => 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 :: (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 = String -> a -> Log Double
forall a. HasCallStack => String -> a
error String
"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 :: 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 :: PriorFunction a
priorFunction = PriorFunction a
prf',
        likelihoodFunction :: PriorFunction a
likelihoodFunction = PriorFunction a
lhf',
        link :: Link a
link = a -> Log Double -> Log Double -> Link a
forall a. a -> Log Double -> Log Double -> Link a
Link a
x (PriorFunction a
prf' a
x) (PriorFunction a
lhf' a
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 :: 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 = String -> IO (MHG a)
forall a. HasCallStack => String -> a
error String
"initMHG: Chain index negative."
  -- Only set the id for the cold chain.
  | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = MHG a -> IO (MHG 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 {chainId :: Maybe Int
chainId = Int -> Maybe Int
forall a. a -> Maybe a
Just Int
0}
  | 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 (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 {chainId :: Maybe Int
chainId = Int -> Maybe Int
forall a. a -> Maybe a
Just Int
i, trace :: Trace a
trace = Trace a
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

-- TODO: Splitmix. Initialization of the MC3 algorithm is an IO action because
-- the generators have to be split. And also because of the mutable trace.

-- | 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 ->
  GenIO ->
  IO (MC3 a)
mc3 :: MC3Settings
-> Settings
-> PriorFunction a
-> PriorFunction a
-> Cycle a
-> Monitor a
-> a
-> GenIO
-> IO (MC3 a)
mc3 MC3Settings
sMc3 Settings
s PriorFunction a
pr PriorFunction a
lh Cycle a
cc Monitor a
mn a
i0 GenIO
g
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 = String -> IO (MC3 a)
forall a. HasCallStack => String -> a
error String
"mc3: The number of chains must be two or larger."
  | Int
sp Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 = String -> IO (MC3 a)
forall a. HasCallStack => String -> a
error String
"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 = String -> IO (MC3 a)
forall a. HasCallStack => String -> a
error String
"mc3: The number of swaps must be in [1, NChains - 1]."
  | Bool
otherwise = do
      -- Split random number generators.
      Vector (Gen RealWorld)
gs <- [Gen RealWorld] -> Vector (Gen RealWorld)
forall a. [a] -> Vector a
V.fromList ([Gen RealWorld] -> Vector (Gen RealWorld))
-> IO [Gen RealWorld] -> IO (Vector (Gen RealWorld))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> GenIO -> IO [GenIO]
forall (m :: * -> *).
PrimMonad m =>
Int -> Gen (PrimState m) -> m [Gen (PrimState m)]
splitGen Int
n GenIO
g
      Vector (MHG a)
cs <- (Gen RealWorld -> IO (MHG a))
-> Vector (Gen RealWorld) -> 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
-> GenIO
-> IO (MHG a)
forall a.
Settings
-> PriorFunction a
-> PriorFunction a
-> Cycle a
-> Monitor a
-> a
-> GenIO
-> IO (MHG a)
mhg Settings
s PriorFunction a
pr PriorFunction a
lh Cycle a
cc Monitor a
mn a
i0) Vector (Gen RealWorld)
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
      MC3 a -> IO (MC3 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
-> Acceptance Int
-> GenIO
-> MC3 a
forall a.
MC3Settings
-> MHGChains a
-> ReciprocalTemperatures
-> Int
-> Acceptance Int
-> GenIO
-> MC3 a
MC3 MC3Settings
sMc3 Vector (MHG a)
hcs ReciprocalTemperatures
bs Int
0 ([Int] -> Acceptance Int
forall k. Ord k => [k] -> Acceptance k
emptyA [Int
0 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2]) GenIO
g
  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 -> String
mc3Fn (AnalysisName String
nm) = String
nm String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
".mcmc.mc3"

-- | Save an MC3 algorithm.
mc3Save ::
  ToJSON a =>
  AnalysisName ->
  MC3 a ->
  IO ()
mc3Save :: 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
  String -> ByteString -> IO ()
BL.writeFile (AnalysisName -> String
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.
--
-- See 'Mcmc.Mcmc.mcmcContinue'.
mc3Load ::
  FromJSON a =>
  PriorFunction a ->
  LikelihoodFunction a ->
  Cycle a ->
  Monitor a ->
  AnalysisName ->
  IO (MC3 a)
mc3Load :: 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 String (SavedMC3 a)
savedMC3 <- ByteString -> Either String (SavedMC3 a)
forall a. FromJSON a => ByteString -> Either String a
eitherDecode (ByteString -> Either String (SavedMC3 a))
-> (ByteString -> ByteString)
-> ByteString
-> Either String (SavedMC3 a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
decompress (ByteString -> Either String (SavedMC3 a))
-> IO ByteString -> IO (Either String (SavedMC3 a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> IO ByteString
BL.readFile (AnalysisName -> String
mc3Fn AnalysisName
nm)
  (String -> IO (MC3 a))
-> (SavedMC3 a -> IO (MC3 a))
-> Either String (SavedMC3 a)
-> IO (MC3 a)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either String -> IO (MC3 a)
forall a. HasCallStack => String -> 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 String (SavedMC3 a)
savedMC3

-- 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 :: 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 = String -> (MHGChains a, Log Double)
forall a. HasCallStack => String -> a
error String
"swapWith: Left index is negative."
  | Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = String -> (MHGChains a, Log Double)
forall a. HasCallStack => String -> a
error String
"swapWith: Right index is negative."
  | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j = String -> (MHGChains a, Log Double)
forall a. HasCallStack => String -> a
error String
"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 :: Link a
link = Link a
ll'}
    cr' :: Chain a
cr' = Chain a
cr {link :: Link a
link = Link a
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 :: 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 -> GenIO -> IO Bool
mhgAccept Log Double
r Gen RealWorld
GenIO
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' :: Acceptance Int
ac' = Int -> Bool -> Acceptance Int -> Acceptance Int
forall k. Ord k => k -> Bool -> Acceptance k -> Acceptance k
pushA Int
i Bool
True (MC3 a -> Acceptance Int
forall a. MC3 a -> Acceptance Int
mc3SwapAcceptance MC3 a
a)
      MC3 a -> IO (MC3 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 :: MHGChains a
mc3MHGChains = MHGChains a
y, mc3SwapAcceptance :: Acceptance Int
mc3SwapAcceptance = Acceptance Int
ac'}
    else do
      let !ac' :: Acceptance Int
ac' = Int -> Bool -> Acceptance Int -> Acceptance Int
forall k. Ord k => k -> Bool -> Acceptance k -> Acceptance k
pushA Int
i Bool
False (MC3 a -> Acceptance Int
forall a. MC3 a -> Acceptance Int
mc3SwapAcceptance MC3 a
a)
      MC3 a -> IO (MC3 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 {mc3SwapAcceptance :: Acceptance Int
mc3SwapAcceptance = Acceptance Int
ac'}
  where
    g :: GenIO
g = MC3 a -> GenIO
forall a. MC3 a -> GenIO
mc3Generator MC3 a
a

mc3IsInValidState :: ToJSON a => MC3 a -> Bool
mc3IsInValidState :: 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

-- TODO: Splitmix. 'mc3Iterate' is actually not parallel, but concurrent because
-- of the IO constraint. Use pure parallel code when we have a pure generator.
--
-- However, we have to take care of the mutable traces.
mc3Iterate ::
  ToJSON a =>
  ParallelizationMode ->
  MC3 a ->
  IO (MC3 a)
mc3Iterate :: ParallelizationMode -> MC3 a -> IO (MC3 a)
mc3Iterate 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] -> GenIO -> IO [Int]
forall a. [a] -> GenIO -> IO [a]
shuffle [Int]
is (GenIO -> IO [Int]) -> GenIO -> IO [Int]
forall a b. (a -> b) -> a -> b
$ MC3 a -> GenIO
forall a. MC3 a -> GenIO
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 (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 (ParallelizationMode -> MHG a -> IO (MHG a)
forall a. Algorithm a => ParallelizationMode -> a -> IO a
aIterate 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 (ParallelizationMode -> MHG a -> IO (MHG a)
forall a. Algorithm a => ParallelizationMode -> a -> IO a
aIterate 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 (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 :: Vector (MHG a)
mc3MHGChains = Vector (MHG a)
mhgs, mc3Iteration :: Int
mc3Iteration = Int -> Int
forall a. Enum a => a -> a
succ Int
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 => Int -> MC3 a -> IO (MC3 a)
mc3AutoTune :: Int -> MC3 a -> IO (MC3 a)
mc3AutoTune 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 (Int -> MHG a -> IO (MHG a)
forall a. Algorithm a => Int -> a -> IO a
aAutoTune 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 = Acceptance Int -> Map Int (Maybe Double)
forall k. Acceptance k -> Map k (Maybe Double)
acceptanceRates (Acceptance Int -> Map Int (Maybe Double))
-> Acceptance Int -> Map Int (Maybe Double)
forall a b. (a -> b) -> a -> b
$ MC3 a -> Acceptance Int
forall a. MC3 a -> Acceptance Int
mc3SwapAcceptance 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 (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 (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 :: Vector (MHG a)
mc3MHGChains = Vector (MHG a)
mhgs'', mc3ReciprocalTemperatures :: ReciprocalTemperatures
mc3ReciprocalTemperatures = ReciprocalTemperatures
bs'}

mc3ResetAcceptance :: ToJSON a => MC3 a -> MC3 a
mc3ResetAcceptance :: MC3 a -> MC3 a
mc3ResetAcceptance 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 MHG a -> MHG a
forall a. Algorithm a => a -> a
aResetAcceptance (MC3 a -> Vector (MHG a)
forall a. MC3 a -> MHGChains a
mc3MHGChains MC3 a
a)
    -- 2. Reset acceptance of swaps.
    ac' :: Acceptance Int
ac' = Acceptance Int -> Acceptance Int
forall k. Ord k => Acceptance k -> Acceptance k
resetA (Acceptance Int -> Acceptance Int)
-> Acceptance Int -> Acceptance Int
forall a b. (a -> b) -> a -> b
$ MC3 a -> Acceptance Int
forall a. MC3 a -> Acceptance Int
mc3SwapAcceptance MC3 a
a
    --
    a' :: MC3 a
a' = MC3 a
a {mc3MHGChains :: Vector (MHG a)
mc3MHGChains = Vector (MHG a)
mhgs', mc3SwapAcceptance :: Acceptance Int
mc3SwapAcceptance = Acceptance Int
ac'}

-- 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 => MC3 a -> BL.ByteString
mc3SummarizeCycle :: MC3 a -> ByteString
mc3SummarizeCycle 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
<> ByteString -> ByteString
BL.fromStrict (Int -> Double -> ByteString
BC.toFixed 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
-> Maybe (Int, Int, Double)
-> ByteString
summarizeProposal
             (String -> PName
PName (String -> PName) -> String -> PName
forall a b. (a -> b) -> a -> b
$ Int -> String
forall a. Show a => a -> String
show Int
i String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" <-> " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
             (String -> PDescription
PDescription String
"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 -> Acceptance Int -> Maybe (Int, Int, Double)
forall k. Ord k => k -> Acceptance k -> Maybe (Int, Int, Double)
acceptanceRate Int
i Acceptance 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 = MHG a -> ByteString
forall a. Algorithm a => a -> ByteString
aSummarizeCycle (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)
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)
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
. Acceptance (Proposal a) -> Map (Proposal a) (Maybe Double)
forall k. Acceptance k -> Map k (Maybe Double)
acceptanceRates (Acceptance (Proposal a) -> Map (Proposal a) (Maybe Double))
-> (Chain a -> Acceptance (Proposal a))
-> Chain a
-> Map (Proposal a) (Maybe Double)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Chain a -> Acceptance (Proposal a)
forall a. Chain a -> Acceptance (Proposal a)
acceptance) 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 (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 (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 (ByteString -> ByteString
BL.fromStrict (ByteString -> ByteString)
-> (Double -> ByteString) -> Double -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Double -> ByteString
BC.toFixed 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 :: Acceptance Int
swapAcceptance = MC3 a -> Acceptance Int
forall a. MC3 a -> Acceptance Int
mc3SwapAcceptance 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

-- No extra monitors are opened.
mc3OpenMonitors :: ToJSON a => AnalysisName -> ExecutionMode -> MC3 a -> IO (MC3 a)
mc3OpenMonitors :: AnalysisName -> ExecutionMode -> MC3 a -> IO (MC3 a)
mc3OpenMonitors AnalysisName
nm ExecutionMode
em 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 (AnalysisName -> ExecutionMode -> MHG a -> IO (MHG a)
forall a. Algorithm a => AnalysisName -> ExecutionMode -> a -> IO a
aOpenMonitors AnalysisName
nm ExecutionMode
em) (MC3 a -> Vector (MHG a)
forall a. MC3 a -> MHGChains a
mc3MHGChains MC3 a
a)
  MC3 a -> IO (MC3 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 :: Vector (MHG a)
mc3MHGChains = Vector (MHG a)
mhgs'}

mc3ExecuteMonitors ::
  ToJSON a =>
  Verbosity ->
  -- Starting time.
  UTCTime ->
  -- Total number of iterations.
  Int ->
  MC3 a ->
  IO (Maybe BL.ByteString)
mc3ExecuteMonitors :: 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 :: 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 :: 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 (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 :: Vector (MHG a)
mc3MHGChains = Vector (MHG a)
mhgs'}