-- |
-- Module      :  Mcmc.Acceptance
-- Description :  Handle acceptance rates
-- Copyright   :  (c) 2021 Dominik Schrempf
-- License     :  GPL-3.0-or-later
--
-- Maintainer  :  dominik.schrempf@gmail.com
-- Stability   :  experimental
-- Portability :  portable
--
-- Creation date: Thu Jul  8 18:12:07 2021.
module Mcmc.Acceptance
  ( -- * Acceptance rates
    AcceptanceRate,
    Acceptance (fromAcceptance),
    emptyA,
    pushA,
    resetA,
    transformKeysA,
    acceptanceRate,
    acceptanceRates,
  )
where

import Control.DeepSeq
import Data.Aeson
import Data.Bifunctor
import Data.Foldable
import qualified Data.Map.Strict as M

-- | Acceptance rate.
type AcceptanceRate = Double

-- | For each key @k@, store the number of accepted and rejected proposals.
newtype Acceptance k = Acceptance {Acceptance k -> Map k (Int, Int)
fromAcceptance :: M.Map k (Int, Int)}
  deriving (Acceptance k -> Acceptance k -> Bool
(Acceptance k -> Acceptance k -> Bool)
-> (Acceptance k -> Acceptance k -> Bool) -> Eq (Acceptance k)
forall k. Eq k => Acceptance k -> Acceptance k -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Acceptance k -> Acceptance k -> Bool
$c/= :: forall k. Eq k => Acceptance k -> Acceptance k -> Bool
== :: Acceptance k -> Acceptance k -> Bool
$c== :: forall k. Eq k => Acceptance k -> Acceptance k -> Bool
Eq, ReadPrec [Acceptance k]
ReadPrec (Acceptance k)
Int -> ReadS (Acceptance k)
ReadS [Acceptance k]
(Int -> ReadS (Acceptance k))
-> ReadS [Acceptance k]
-> ReadPrec (Acceptance k)
-> ReadPrec [Acceptance k]
-> Read (Acceptance k)
forall k. (Ord k, Read k) => ReadPrec [Acceptance k]
forall k. (Ord k, Read k) => ReadPrec (Acceptance k)
forall k. (Ord k, Read k) => Int -> ReadS (Acceptance k)
forall k. (Ord k, Read k) => ReadS [Acceptance k]
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
readListPrec :: ReadPrec [Acceptance k]
$creadListPrec :: forall k. (Ord k, Read k) => ReadPrec [Acceptance k]
readPrec :: ReadPrec (Acceptance k)
$creadPrec :: forall k. (Ord k, Read k) => ReadPrec (Acceptance k)
readList :: ReadS [Acceptance k]
$creadList :: forall k. (Ord k, Read k) => ReadS [Acceptance k]
readsPrec :: Int -> ReadS (Acceptance k)
$creadsPrec :: forall k. (Ord k, Read k) => Int -> ReadS (Acceptance k)
Read, Int -> Acceptance k -> ShowS
[Acceptance k] -> ShowS
Acceptance k -> String
(Int -> Acceptance k -> ShowS)
-> (Acceptance k -> String)
-> ([Acceptance k] -> ShowS)
-> Show (Acceptance k)
forall k. Show k => Int -> Acceptance k -> ShowS
forall k. Show k => [Acceptance k] -> ShowS
forall k. Show k => Acceptance k -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Acceptance k] -> ShowS
$cshowList :: forall k. Show k => [Acceptance k] -> ShowS
show :: Acceptance k -> String
$cshow :: forall k. Show k => Acceptance k -> String
showsPrec :: Int -> Acceptance k -> ShowS
$cshowsPrec :: forall k. Show k => Int -> Acceptance k -> ShowS
Show)

instance ToJSONKey k => ToJSON (Acceptance k) where
  toJSON :: Acceptance k -> Value
toJSON (Acceptance Map k (Int, Int)
m) = Map k (Int, Int) -> Value
forall a. ToJSON a => a -> Value
toJSON Map k (Int, Int)
m
  toEncoding :: Acceptance k -> Encoding
toEncoding (Acceptance Map k (Int, Int)
m) = Map k (Int, Int) -> Encoding
forall a. ToJSON a => a -> Encoding
toEncoding Map k (Int, Int)
m

instance (Ord k, FromJSONKey k) => FromJSON (Acceptance k) where
  parseJSON :: Value -> Parser (Acceptance k)
parseJSON Value
v = Map k (Int, Int) -> Acceptance k
forall k. Map k (Int, Int) -> Acceptance k
Acceptance (Map k (Int, Int) -> Acceptance k)
-> Parser (Map k (Int, Int)) -> Parser (Acceptance k)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Value -> Parser (Map k (Int, Int))
forall a. FromJSON a => Value -> Parser a
parseJSON Value
v

-- | In the beginning there was the Word.
--
-- Initialize an empty storage of accepted/rejected values.
emptyA :: Ord k => [k] -> Acceptance k
emptyA :: [k] -> Acceptance k
emptyA [k]
ks = Map k (Int, Int) -> Acceptance k
forall k. Map k (Int, Int) -> Acceptance k
Acceptance (Map k (Int, Int) -> Acceptance k)
-> Map k (Int, Int) -> Acceptance k
forall a b. (a -> b) -> a -> b
$ [(k, (Int, Int))] -> Map k (Int, Int)
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(k
k, (Int
0, Int
0)) | k
k <- [k]
ks]

-- | For key @k@, prepend an accepted (True) or rejected (False) proposal.
pushA :: Ord k => k -> Bool -> Acceptance k -> Acceptance k
pushA :: k -> Bool -> Acceptance k -> Acceptance k
pushA k
k Bool
True = Map k (Int, Int) -> Acceptance k
forall k. Map k (Int, Int) -> Acceptance k
Acceptance (Map k (Int, Int) -> Acceptance k)
-> (Acceptance k -> Map k (Int, Int))
-> Acceptance k
-> Acceptance k
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Int, Int) -> (Int, Int))
-> k -> Map k (Int, Int) -> Map k (Int, Int)
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust ((Int, Int) -> (Int, Int)
forall a. NFData a => a -> a
force ((Int, Int) -> (Int, Int))
-> ((Int, Int) -> (Int, Int)) -> (Int, Int) -> (Int, Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int) -> (Int, Int) -> (Int, Int)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Int -> Int
forall a. Enum a => a -> a
succ) k
k (Map k (Int, Int) -> Map k (Int, Int))
-> (Acceptance k -> Map k (Int, Int))
-> Acceptance k
-> Map k (Int, Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Acceptance k -> Map k (Int, Int)
forall k. Acceptance k -> Map k (Int, Int)
fromAcceptance
pushA k
k Bool
False = Map k (Int, Int) -> Acceptance k
forall k. Map k (Int, Int) -> Acceptance k
Acceptance (Map k (Int, Int) -> Acceptance k)
-> (Acceptance k -> Map k (Int, Int))
-> Acceptance k
-> Acceptance k
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Int, Int) -> (Int, Int))
-> k -> Map k (Int, Int) -> Map k (Int, Int)
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust ((Int, Int) -> (Int, Int)
forall a. NFData a => a -> a
force ((Int, Int) -> (Int, Int))
-> ((Int, Int) -> (Int, Int)) -> (Int, Int) -> (Int, Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int) -> (Int, Int) -> (Int, Int)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second Int -> Int
forall a. Enum a => a -> a
succ) k
k (Map k (Int, Int) -> Map k (Int, Int))
-> (Acceptance k -> Map k (Int, Int))
-> Acceptance k
-> Map k (Int, Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Acceptance k -> Map k (Int, Int)
forall k. Acceptance k -> Map k (Int, Int)
fromAcceptance

-- | Reset acceptance storage.
resetA :: Ord k => Acceptance k -> Acceptance k
resetA :: Acceptance k -> Acceptance k
resetA = [k] -> Acceptance k
forall k. Ord k => [k] -> Acceptance k
emptyA ([k] -> Acceptance k)
-> (Acceptance k -> [k]) -> Acceptance k -> Acceptance k
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map k (Int, Int) -> [k]
forall k a. Map k a -> [k]
M.keys (Map k (Int, Int) -> [k])
-> (Acceptance k -> Map k (Int, Int)) -> Acceptance k -> [k]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Acceptance k -> Map k (Int, Int)
forall k. Acceptance k -> Map k (Int, Int)
fromAcceptance

transformKeys :: (Ord k1, Ord k2) => [k1] -> [k2] -> M.Map k1 v -> M.Map k2 v
transformKeys :: [k1] -> [k2] -> Map k1 v -> Map k2 v
transformKeys [k1]
ks1 [k2]
ks2 Map k1 v
m = (Map k2 v -> (k1, k2) -> Map k2 v)
-> Map k2 v -> [(k1, k2)] -> Map k2 v
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Map k2 v -> (k1, k2) -> Map k2 v
forall k. Ord k => Map k v -> (k1, k) -> Map k v
insrt Map k2 v
forall k a. Map k a
M.empty ([(k1, k2)] -> Map k2 v) -> [(k1, k2)] -> Map k2 v
forall a b. (a -> b) -> a -> b
$ [k1] -> [k2] -> [(k1, k2)]
forall a b. [a] -> [b] -> [(a, b)]
zip [k1]
ks1 [k2]
ks2
  where
    insrt :: Map k v -> (k1, k) -> Map k v
insrt Map k v
m' (k1
k1, k
k2) = k -> v -> Map k v -> Map k v
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
k2 (Map k1 v
m Map k1 v -> k1 -> v
forall k a. Ord k => Map k a -> k -> a
M.! k1
k1) Map k v
m'

-- | Transform keys using the given lists. Keys not provided will not be present
-- in the new 'Acceptance' variable.
transformKeysA :: (Ord k1, Ord k2) => [k1] -> [k2] -> Acceptance k1 -> Acceptance k2
transformKeysA :: [k1] -> [k2] -> Acceptance k1 -> Acceptance k2
transformKeysA [k1]
ks1 [k2]
ks2 = Map k2 (Int, Int) -> Acceptance k2
forall k. Map k (Int, Int) -> Acceptance k
Acceptance (Map k2 (Int, Int) -> Acceptance k2)
-> (Acceptance k1 -> Map k2 (Int, Int))
-> Acceptance k1
-> Acceptance k2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [k1] -> [k2] -> Map k1 (Int, Int) -> Map k2 (Int, Int)
forall k1 k2 v.
(Ord k1, Ord k2) =>
[k1] -> [k2] -> Map k1 v -> Map k2 v
transformKeys [k1]
ks1 [k2]
ks2 (Map k1 (Int, Int) -> Map k2 (Int, Int))
-> (Acceptance k1 -> Map k1 (Int, Int))
-> Acceptance k1
-> Map k2 (Int, Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Acceptance k1 -> Map k1 (Int, Int)
forall k. Acceptance k -> Map k (Int, Int)
fromAcceptance

-- | Acceptance counts and rate for a specific proposal.
--
-- Return 'Nothing' if no proposals have been accepted or rejected (division by
-- zero).
acceptanceRate :: Ord k => k -> Acceptance k -> Maybe (Int, Int, AcceptanceRate)
acceptanceRate :: k -> Acceptance k -> Maybe (Int, Int, AcceptanceRate)
acceptanceRate k
k Acceptance k
a = case Acceptance k -> Map k (Int, Int)
forall k. Acceptance k -> Map k (Int, Int)
fromAcceptance Acceptance k
a Map k (Int, Int) -> k -> Maybe (Int, Int)
forall k a. Ord k => Map k a -> k -> Maybe a
M.!? k
k of
  Just (Int
0, Int
0) -> Maybe (Int, Int, AcceptanceRate)
forall a. Maybe a
Nothing
  Just (Int
as, Int
rs) -> (Int, Int, AcceptanceRate) -> Maybe (Int, Int, AcceptanceRate)
forall a. a -> Maybe a
Just (Int
as, Int
rs, Int -> AcceptanceRate
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
as AcceptanceRate -> AcceptanceRate -> AcceptanceRate
forall a. Fractional a => a -> a -> a
/ Int -> AcceptanceRate
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
as Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
rs))
  Maybe (Int, Int)
Nothing -> String -> Maybe (Int, Int, AcceptanceRate)
forall a. HasCallStack => String -> a
error String
"acceptanceRate: Key not found in map."

-- | Acceptance rates for all proposals.
--
-- Set rate to 'Nothing' if no proposals have been accepted or rejected
-- (division by zero).
acceptanceRates :: Acceptance k -> M.Map k (Maybe AcceptanceRate)
acceptanceRates :: Acceptance k -> Map k (Maybe AcceptanceRate)
acceptanceRates =
  ((Int, Int) -> Maybe AcceptanceRate)
-> Map k (Int, Int) -> Map k (Maybe AcceptanceRate)
forall a b k. (a -> b) -> Map k a -> Map k b
M.map
    ( \(Int
as, Int
rs) ->
        if Int
as Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
rs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
          then Maybe AcceptanceRate
forall a. Maybe a
Nothing
          else AcceptanceRate -> Maybe AcceptanceRate
forall a. a -> Maybe a
Just (AcceptanceRate -> Maybe AcceptanceRate)
-> AcceptanceRate -> Maybe AcceptanceRate
forall a b. (a -> b) -> a -> b
$ Int -> AcceptanceRate
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
as AcceptanceRate -> AcceptanceRate -> AcceptanceRate
forall a. Fractional a => a -> a -> a
/ Int -> AcceptanceRate
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
as Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
rs)
    )
    (Map k (Int, Int) -> Map k (Maybe AcceptanceRate))
-> (Acceptance k -> Map k (Int, Int))
-> Acceptance k
-> Map k (Maybe AcceptanceRate)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Acceptance k -> Map k (Int, Int)
forall k. Acceptance k -> Map k (Int, Int)
fromAcceptance