{-# LANGUAGE TemplateHaskell #-}

-- |
-- Module      :  Mcmc.Acceptance
-- Description :  Handle acceptance rates
-- Copyright   :  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,
    AcceptanceCounts (..),
    Acceptance (fromAcceptance),
    emptyA,
    pushAccept,
    pushReject,
    pushAcceptanceCounts,
    resetA,
    transformKeysA,
    acceptanceRate,
    acceptanceRates,
  )
where

import Data.Aeson
import Data.Aeson.TH
import Data.Foldable
import qualified Data.Map.Strict as M

-- | Acceptance rate.
type AcceptanceRate = Double

-- | Number of accepted and rejected proposals.
data AcceptanceCounts = AcceptanceCounts
  { AcceptanceCounts -> Int
nAccepted :: !Int,
    AcceptanceCounts -> Int
nRejected :: !Int
  }
  deriving (Int -> AcceptanceCounts -> ShowS
[AcceptanceCounts] -> ShowS
AcceptanceCounts -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [AcceptanceCounts] -> ShowS
$cshowList :: [AcceptanceCounts] -> ShowS
show :: AcceptanceCounts -> String
$cshow :: AcceptanceCounts -> String
showsPrec :: Int -> AcceptanceCounts -> ShowS
$cshowsPrec :: Int -> AcceptanceCounts -> ShowS
Show, AcceptanceCounts -> AcceptanceCounts -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: AcceptanceCounts -> AcceptanceCounts -> Bool
$c/= :: AcceptanceCounts -> AcceptanceCounts -> Bool
== :: AcceptanceCounts -> AcceptanceCounts -> Bool
$c== :: AcceptanceCounts -> AcceptanceCounts -> Bool
Eq, Eq AcceptanceCounts
AcceptanceCounts -> AcceptanceCounts -> Bool
AcceptanceCounts -> AcceptanceCounts -> Ordering
AcceptanceCounts -> AcceptanceCounts -> AcceptanceCounts
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: AcceptanceCounts -> AcceptanceCounts -> AcceptanceCounts
$cmin :: AcceptanceCounts -> AcceptanceCounts -> AcceptanceCounts
max :: AcceptanceCounts -> AcceptanceCounts -> AcceptanceCounts
$cmax :: AcceptanceCounts -> AcceptanceCounts -> AcceptanceCounts
>= :: AcceptanceCounts -> AcceptanceCounts -> Bool
$c>= :: AcceptanceCounts -> AcceptanceCounts -> Bool
> :: AcceptanceCounts -> AcceptanceCounts -> Bool
$c> :: AcceptanceCounts -> AcceptanceCounts -> Bool
<= :: AcceptanceCounts -> AcceptanceCounts -> Bool
$c<= :: AcceptanceCounts -> AcceptanceCounts -> Bool
< :: AcceptanceCounts -> AcceptanceCounts -> Bool
$c< :: AcceptanceCounts -> AcceptanceCounts -> Bool
compare :: AcceptanceCounts -> AcceptanceCounts -> Ordering
$ccompare :: AcceptanceCounts -> AcceptanceCounts -> Ordering
Ord)

$(deriveJSON defaultOptions ''AcceptanceCounts)

addAccept :: AcceptanceCounts -> AcceptanceCounts
addAccept :: AcceptanceCounts -> AcceptanceCounts
addAccept (AcceptanceCounts Int
a Int
r) = Int -> Int -> AcceptanceCounts
AcceptanceCounts (Int
a forall a. Num a => a -> a -> a
+ Int
1) Int
r

addReject :: AcceptanceCounts -> AcceptanceCounts
addReject :: AcceptanceCounts -> AcceptanceCounts
addReject (AcceptanceCounts Int
a Int
r) = Int -> Int -> AcceptanceCounts
AcceptanceCounts Int
a (Int
r forall a. Num a => a -> a -> a
+ Int
1)

addAcceptanceCounts :: AcceptanceCounts -> AcceptanceCounts -> AcceptanceCounts
addAcceptanceCounts :: AcceptanceCounts -> AcceptanceCounts -> AcceptanceCounts
addAcceptanceCounts (AcceptanceCounts Int
al Int
rl) (AcceptanceCounts Int
ar Int
rr) =
  Int -> Int -> AcceptanceCounts
AcceptanceCounts (Int
al forall a. Num a => a -> a -> a
+ Int
ar) (Int
rl forall a. Num a => a -> a -> a
+ Int
rr)

-- | For each key @k@, store the number of accepted and rejected proposals.
newtype Acceptance k = Acceptance {forall k. Acceptance k -> Map k AcceptanceCounts
fromAcceptance :: M.Map k AcceptanceCounts}
  deriving (Acceptance k -> Acceptance k -> Bool
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, Int -> Acceptance k -> ShowS
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 AcceptanceCounts
m) = forall a. ToJSON a => a -> Value
toJSON Map k AcceptanceCounts
m
  toEncoding :: Acceptance k -> Encoding
toEncoding (Acceptance Map k AcceptanceCounts
m) = forall a. ToJSON a => a -> Encoding
toEncoding Map k AcceptanceCounts
m

instance (Ord k, FromJSONKey k) => FromJSON (Acceptance k) where
  parseJSON :: Value -> Parser (Acceptance k)
parseJSON Value
v = forall k. Map k AcceptanceCounts -> Acceptance k
Acceptance forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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 :: forall k. Ord k => [k] -> Acceptance k
emptyA [k]
ks = forall k. Map k AcceptanceCounts -> Acceptance k
Acceptance forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(k
k, Int -> Int -> AcceptanceCounts
AcceptanceCounts Int
0 Int
0) | k
k <- [k]
ks]

-- | For key @k@, add an accept.
pushAccept :: Ord k => k -> Acceptance k -> Acceptance k
pushAccept :: forall k. Ord k => k -> Acceptance k -> Acceptance k
pushAccept k
k = forall k. Map k AcceptanceCounts -> Acceptance k
Acceptance forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust AcceptanceCounts -> AcceptanceCounts
addAccept k
k forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k. Acceptance k -> Map k AcceptanceCounts
fromAcceptance

-- | For key @k@, add a reject.
pushReject :: Ord k => k -> Acceptance k -> Acceptance k
pushReject :: forall k. Ord k => k -> Acceptance k -> Acceptance k
pushReject k
k = forall k. Map k AcceptanceCounts -> Acceptance k
Acceptance forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust AcceptanceCounts -> AcceptanceCounts
addReject k
k forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k. Acceptance k -> Map k AcceptanceCounts
fromAcceptance

-- | For key @k@, add acceptance counts.
pushAcceptanceCounts :: Ord k => k -> AcceptanceCounts -> Acceptance k -> Acceptance k
pushAcceptanceCounts :: forall k.
Ord k =>
k -> AcceptanceCounts -> Acceptance k -> Acceptance k
pushAcceptanceCounts k
k AcceptanceCounts
c = forall k. Map k AcceptanceCounts -> Acceptance k
Acceptance forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust (AcceptanceCounts -> AcceptanceCounts -> AcceptanceCounts
addAcceptanceCounts AcceptanceCounts
c) k
k forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k. Acceptance k -> Map k AcceptanceCounts
fromAcceptance

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

transformKeys :: (Ord k1, Ord k2) => [(k1, k2)] -> M.Map k1 v -> M.Map k2 v
transformKeys :: forall k1 k2 v.
(Ord k1, Ord k2) =>
[(k1, k2)] -> Map k1 v -> Map k2 v
transformKeys [(k1, k2)]
ks Map k1 v
m = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall {k}. Ord k => Map k v -> (k1, k) -> Map k v
insrt forall k a. Map k a
M.empty [(k1, k2)]
ks
  where
    insrt :: Map k v -> (k1, k) -> Map k v
insrt Map k v
m' (k1
k1, k
k2) = forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
k2 (Map k1 v
m 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 :: forall k1 k2.
(Ord k1, Ord k2) =>
[(k1, k2)] -> Acceptance k1 -> Acceptance k2
transformKeysA [(k1, k2)]
ks = forall k. Map k AcceptanceCounts -> Acceptance k
Acceptance forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k1 k2 v.
(Ord k1, Ord k2) =>
[(k1, k2)] -> Map k1 v -> Map k2 v
transformKeys [(k1, k2)]
ks forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k. Acceptance k -> Map k AcceptanceCounts
fromAcceptance

-- | Acceptance counts and rate for a specific proposal.
--
-- Return @Just (accepts, rejects, acceptance rate)@.
--
-- Return 'Nothing' if no proposals have been accepted or rejected (division by
-- zero).
acceptanceRate ::
  Ord k =>
  k ->
  Acceptance k ->
  Maybe (Int, Int, AcceptanceRate)
acceptanceRate :: forall k.
Ord k =>
k -> Acceptance k -> Maybe (Int, Int, AcceptanceRate)
acceptanceRate k
k Acceptance k
a = case forall k. Acceptance k -> Map k AcceptanceCounts
fromAcceptance Acceptance k
a forall k a. Ord k => Map k a -> k -> Maybe a
M.!? k
k of
  Just (AcceptanceCounts Int
0 Int
0) -> forall a. Maybe a
Nothing
  Just (AcceptanceCounts Int
as Int
rs) -> forall a. a -> Maybe a
Just (Int
as, Int
rs, forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
as forall a. Fractional a => a -> a -> a
/ forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
as forall a. Num a => a -> a -> a
+ Int
rs))
  Maybe AcceptanceCounts
Nothing -> 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 :: forall k. Acceptance k -> Map k (Maybe AcceptanceRate)
acceptanceRates =
  forall a b k. (a -> b) -> Map k a -> Map k b
M.map
    ( \(AcceptanceCounts Int
as Int
rs) ->
        if Int
as forall a. Num a => a -> a -> a
+ Int
rs forall a. Eq a => a -> a -> Bool
== Int
0
          then forall a. Maybe a
Nothing
          else forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
as forall a. Fractional a => a -> a -> a
/ forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
as forall a. Num a => a -> a -> a
+ Int
rs)
    )
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k. Acceptance k -> Map k AcceptanceCounts
fromAcceptance