{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskell #-}
module Mcmc.Chain.Save
( SavedChain (..),
toSavedChain,
fromSavedChain,
fromSavedChainUnsafe,
)
where
import Control.Monad
import Data.Aeson
import Data.Aeson.TH
import Data.List hiding (cycle)
import Data.Maybe
import qualified Data.Stack.Circular as C
import qualified Data.Vector as VB
import Data.Word
import Mcmc.Acceptance
import Mcmc.Chain.Chain
import Mcmc.Chain.Link
import Mcmc.Chain.Trace
import Mcmc.Cycle
import Mcmc.Internal.Random
import Mcmc.Likelihood
import Mcmc.Monitor
import Mcmc.Prior
import Mcmc.Proposal
import Prelude hiding (cycle)
data SavedChain a = SavedChain
{ forall a. SavedChain a -> Link a
savedLink :: Link a,
forall a. SavedChain a -> Int
savedIteration :: Int,
forall a. SavedChain a -> Stack Vector (Link a)
savedTrace :: C.Stack VB.Vector (Link a),
forall a. SavedChain a -> Acceptances Int
savedAcceptances :: Acceptances Int,
forall a. SavedChain a -> (Word64, Word64)
savedSeed :: (Word64, Word64),
forall a.
SavedChain a
-> [Maybe (TuningParameter, AuxiliaryTuningParameters)]
savedTuningParameters :: [Maybe (TuningParameter, AuxiliaryTuningParameters)]
}
deriving (SavedChain a -> SavedChain a -> Bool
forall a. Eq a => SavedChain a -> SavedChain a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: SavedChain a -> SavedChain a -> Bool
$c/= :: forall a. Eq a => SavedChain a -> SavedChain a -> Bool
== :: SavedChain a -> SavedChain a -> Bool
$c== :: forall a. Eq a => SavedChain a -> SavedChain a -> Bool
Eq, Int -> SavedChain a -> ShowS
forall a. Show a => Int -> SavedChain a -> ShowS
forall a. Show a => [SavedChain a] -> ShowS
forall a. Show a => SavedChain a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [SavedChain a] -> ShowS
$cshowList :: forall a. Show a => [SavedChain a] -> ShowS
show :: SavedChain a -> String
$cshow :: forall a. Show a => SavedChain a -> String
showsPrec :: Int -> SavedChain a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> SavedChain a -> ShowS
Show)
$(deriveJSON defaultOptions ''SavedChain)
toSavedChain ::
Chain a ->
IO (SavedChain a)
toSavedChain :: forall a. Chain a -> IO (SavedChain a)
toSavedChain (Chain Link a
it Int
i Trace a
tr Acceptances (Proposal a)
ac IOGenM StdGen
g Int
_ PriorFunction a
_ PriorFunction a
_ Cycle a
cc Monitor a
_) = do
(Word64, Word64)
g' <- IOGenM StdGen -> IO (Word64, Word64)
saveGen IOGenM StdGen
g
Stack Vector (Link a)
tr' <- forall a. Trace a -> IO (Stack Vector (Link a))
freezeT Trace a
tr
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a.
Link a
-> Int
-> Stack Vector (Link a)
-> Acceptances Int
-> (Word64, Word64)
-> [Maybe (TuningParameter, AuxiliaryTuningParameters)]
-> SavedChain a
SavedChain Link a
it Int
i Stack Vector (Link a)
tr' Acceptances Int
ac' (Word64, Word64)
g' [Maybe (TuningParameter, AuxiliaryTuningParameters)]
ts
where
ps :: [Proposal a]
ps = forall a. Cycle a -> [Proposal a]
ccProposals Cycle a
cc
ac' :: Acceptances Int
ac' = forall k1 k2.
(Ord k1, Ord k2) =>
[(k1, k2)] -> Acceptances k1 -> Acceptances k2
transformKeysA (forall a b. [a] -> [b] -> [(a, b)]
zip [Proposal a]
ps [Int
0 ..]) Acceptances (Proposal a)
ac
ts :: [Maybe (TuningParameter, AuxiliaryTuningParameters)]
ts =
[ (\Tuner a
t -> (forall a. Tuner a -> TuningParameter
tTuningParameter Tuner a
t, forall a. Tuner a -> AuxiliaryTuningParameters
tAuxiliaryTuningParameters Tuner a
t)) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (Tuner a)
mt
| Maybe (Tuner a)
mt <- forall a b. (a -> b) -> [a] -> [b]
map forall a. Proposal a -> Maybe (Tuner a)
prTuner [Proposal a]
ps
]
fromSavedChain ::
PriorFunction a ->
LikelihoodFunction a ->
Cycle a ->
Monitor a ->
SavedChain a ->
IO (Chain a)
fromSavedChain :: forall a.
PriorFunction a
-> PriorFunction a
-> Cycle a
-> Monitor a
-> SavedChain a
-> IO (Chain a)
fromSavedChain PriorFunction a
pr PriorFunction a
lh Cycle a
cc Monitor a
mn SavedChain a
sv
| PriorFunction a
pr (forall a. Link a -> a
state Link a
it) forall a. Eq a => a -> a -> Bool
/= forall a. Link a -> Log TuningParameter
prior Link a
it =
let msg :: String
msg =
[String] -> String
unlines
[ String
"fromSave: Provided prior function does not match the saved prior.",
String
"fromSave: Current prior:" forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show (forall a. Link a -> Log TuningParameter
prior Link a
it) forall a. Semigroup a => a -> a -> a
<> String
".",
String
"fromSave: Given prior:" forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show (PriorFunction a
pr forall a b. (a -> b) -> a -> b
$ forall a. Link a -> a
state Link a
it) forall a. Semigroup a => a -> a -> a
<> String
"."
]
in forall a. HasCallStack => String -> a
error String
msg
| PriorFunction a
lh (forall a. Link a -> a
state Link a
it) forall a. Eq a => a -> a -> Bool
/= forall a. Link a -> Log TuningParameter
likelihood Link a
it =
let msg :: String
msg =
[String] -> String
unlines
[ String
"fromSave: Provided likelihood function does not match the saved likelihood function.",
String
"fromSave: Current likelihood:" forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show (forall a. Link a -> Log TuningParameter
likelihood Link a
it) forall a. Semigroup a => a -> a -> a
<> String
".",
String
"fromSave: Given likelihood:" forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show (PriorFunction a
lh forall a b. (a -> b) -> a -> b
$ forall a. Link a -> a
state Link a
it) forall a. Semigroup a => a -> a -> a
<> String
"."
]
in forall a. HasCallStack => String -> a
error String
msg
| forall (t :: * -> *) a. Foldable t => t a -> Int
length (forall k. Acceptances k -> Map k Acceptance
fromAcceptances Acceptances Int
ac) forall a. Eq a => a -> a -> Bool
/= forall (t :: * -> *) a. Foldable t => t a -> Int
length (forall a. Cycle a -> [Proposal a]
ccProposals Cycle a
cc) =
let msg :: String
msg =
[String] -> String
unlines
[ String
"fromSave: The number of proposals does not match.",
String
"fromSave: Number of saved proposals:" forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show (forall (t :: * -> *) a. Foldable t => t a -> Int
length forall a b. (a -> b) -> a -> b
$ forall k. Acceptances k -> Map k Acceptance
fromAcceptances Acceptances Int
ac) forall a. Semigroup a => a -> a -> a
<> String
".",
String
"fromSave: Number of given proposals:" forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show (forall (t :: * -> *) a. Foldable t => t a -> Int
length forall a b. (a -> b) -> a -> b
$ forall a. Cycle a -> [Proposal a]
ccProposals Cycle a
cc) forall a. Semigroup a => a -> a -> a
<> String
"."
]
in forall a. HasCallStack => String -> a
error String
msg
| Bool
otherwise = forall a.
PriorFunction a
-> PriorFunction a
-> Cycle a
-> Monitor a
-> SavedChain a
-> IO (Chain a)
fromSavedChainUnsafe PriorFunction a
pr PriorFunction a
lh Cycle a
cc Monitor a
mn SavedChain a
sv
where
it :: Link a
it = forall a. SavedChain a -> Link a
savedLink SavedChain a
sv
ac :: Acceptances Int
ac = forall a. SavedChain a -> Acceptances Int
savedAcceptances SavedChain a
sv
fromSavedChainUnsafe ::
PriorFunction a ->
LikelihoodFunction a ->
Cycle a ->
Monitor a ->
SavedChain a ->
IO (Chain a)
fromSavedChainUnsafe :: forall a.
PriorFunction a
-> PriorFunction a
-> Cycle a
-> Monitor a
-> SavedChain a
-> IO (Chain a)
fromSavedChainUnsafe PriorFunction a
pr PriorFunction a
lh Cycle a
cc Monitor a
mn (SavedChain Link a
it Int
i Stack Vector (Link a)
tr Acceptances Int
ac' (Word64, Word64)
g' [Maybe (TuningParameter, AuxiliaryTuningParameters)]
ts) = do
IOGenM StdGen
g <- (Word64, Word64) -> IO (IOGenM StdGen)
loadGen (Word64, Word64)
g'
Trace a
tr' <- forall a. Stack Vector (Link a) -> IO (Trace a)
thawT Stack Vector (Link a)
tr
forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a.
Link a
-> Int
-> Trace a
-> Acceptances (Proposal a)
-> IOGenM StdGen
-> Int
-> PriorFunction a
-> PriorFunction a
-> Cycle a
-> Monitor a
-> Chain a
Chain Link a
it Int
i Trace a
tr' Acceptances (Proposal a)
ac IOGenM StdGen
g Int
i PriorFunction a
pr PriorFunction a
lh Cycle a
cc' Monitor a
mn
where
ac :: Acceptances (Proposal a)
ac = forall k1 k2.
(Ord k1, Ord k2) =>
[(k1, k2)] -> Acceptances k1 -> Acceptances k2
transformKeysA (forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] (forall a. Cycle a -> [Proposal a]
ccProposals Cycle a
cc)) Acceptances Int
ac'
tunePs :: Maybe (TuningParameter, AuxiliaryTuningParameters)
-> Proposal a -> Proposal a
tunePs Maybe (TuningParameter, AuxiliaryTuningParameters)
mt Proposal a
p = case Maybe (TuningParameter, AuxiliaryTuningParameters)
mt of
Maybe (TuningParameter, AuxiliaryTuningParameters)
Nothing -> Proposal a
p
Just (TuningParameter
x, AuxiliaryTuningParameters
xs) -> forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (forall a. HasCallStack => String -> a
error forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a. Semigroup a => a -> a -> a
<> forall {a}. a
err)) forall a. a -> a
id forall a b. (a -> b) -> a -> b
$ forall a.
TuningParameter
-> AuxiliaryTuningParameters
-> Proposal a
-> Either String (Proposal a)
tuneWithTuningParameters TuningParameter
x AuxiliaryTuningParameters
xs Proposal a
p
err :: a
err = forall a. HasCallStack => String -> a
error String
"\nfromSavedChain: Proposal with stored tuning parameters is not tunable."
ps :: [Proposal a]
ps = forall a. Cycle a -> [Proposal a]
ccProposals Cycle a
cc
cc' :: Cycle a
cc' = Cycle a
cc {ccProposals :: [Proposal a]
ccProposals = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall {a}.
Maybe (TuningParameter, AuxiliaryTuningParameters)
-> Proposal a -> Proposal a
tunePs [Maybe (TuningParameter, AuxiliaryTuningParameters)]
ts [Proposal a]
ps}