-- |
-- Module      :  Mcmc.Proposal.Generic
-- Description :  Generic interface for creating proposals
-- Copyright   :  (c) Dominik Schrempf 2021
-- License     :  GPL-3.0-or-later
--
-- Maintainer  :  dominik.schrempf@gmail.com
-- Stability   :  unstable
-- Portability :  portable
--
-- Creation date: Thu May 14 20:26:27 2020.
module Mcmc.Proposal.Generic
  ( genericContinuous,
    genericDiscrete,
  )
where

import Mcmc.Proposal
import Numeric.Log
import Statistics.Distribution

-- | Generic function to create proposals for continuous parameters (e.g.,
-- 'Double').
--
-- The procedure is as follows: Let \(\mathbb{X}\) be the state space and \(x\)
-- be the current state.
--
-- 1. Let \(D\) be a continuous probability distribution on \(\mathbb{D}\);
--    sample an auxiliary variable \(epsilon \sim D\).
--
-- 2. Suppose \(\odot : \mathbb{X} \times \mathbb{D} \to \mathbb{X}\). Propose a
--    new state \(x' = x \odot \epsilon\).
--
-- 3. If the proposal is unbiased, the Metropolis-Hastings-Green ratio can
--    directly be calculated using the posterior function.
--
-- 4. However, if the proposal is biased: Suppose \(g : \mathbb{D} \to
--    \mathbb{D}\) inverses the auxiliary variable \(\epsilon\) such that \(x =
--    x' \odot g(\epsilon)\). Calculate the Metropolis-Hastings-Green ratio
--    using the posterior function, \(g\), \(D\), \(\epsilon\), and possibly a
--    Jacobian function.
genericContinuous ::
  (ContDistr d, ContGen d) =>
  -- | Probability distribution
  d ->
  -- | Forward operator \(\odot\).
  --
  -- For example, for a multiplicative proposal on one variable the forward
  -- operator is @(*)@, so that @x * u = y@.
  (a -> Double -> a) ->
  -- | Inverse operator \(g\) of the auxiliary variable.
  --
  -- For example, 'recip' for a multiplicative proposal on one variable, since
  -- @y * (recip u) = x * u * (recip u) = x@.
  --
  -- Required for biased proposals.
  Maybe (Double -> Double) ->
  -- | Function to compute the absolute value of the determinant of the Jacobian
  -- matrix. For example, for a multiplicative proposal on one variable, we have
  --
  -- @
  -- detJacobian _ u = Exp $ log $ recip u
  -- @
  --
  -- That is, the determinant of the Jacobian matrix of multiplication is just
  -- the reciprocal value of @u@ (with conversion to log domain).
  --
  -- Required for proposals for which the absolute value of the determinant of
  -- the Jacobian differs from 1.0.
  --
  -- Conversion to log domain is necessary, because some determinants of
  -- Jacobians are very small (or large).
  Maybe (a -> Double -> Jacobian) ->
  ProposalSimple a
genericContinuous :: d
-> (a -> Double -> a)
-> Maybe (Double -> Double)
-> Maybe (a -> Double -> Jacobian)
-> ProposalSimple a
genericContinuous d
d a -> Double -> a
f Maybe (Double -> Double)
mInv Maybe (a -> Double -> Jacobian)
mJac a
x GenIO
g = do
  Double
u <- d -> Gen RealWorld -> IO Double
forall d g (m :: * -> *).
(ContGen d, StatefulGen g m) =>
d -> g -> m Double
genContVar d
d Gen RealWorld
GenIO
g
  let r :: Jacobian
r = case Maybe (Double -> Double)
mInv of
        Maybe (Double -> Double)
Nothing -> Jacobian
1.0
        Just Double -> Double
fInv ->
          let qXY :: Jacobian
qXY = Double -> Jacobian
forall a. a -> Log a
Exp (Double -> Jacobian) -> Double -> Jacobian
forall a b. (a -> b) -> a -> b
$ d -> Double -> Double
forall d. ContDistr d => d -> Double -> Double
logDensity d
d Double
u
              qYX :: Jacobian
qYX = Double -> Jacobian
forall a. a -> Log a
Exp (Double -> Jacobian) -> Double -> Jacobian
forall a b. (a -> b) -> a -> b
$ d -> Double -> Double
forall d. ContDistr d => d -> Double -> Double
logDensity d
d (Double -> Double
fInv Double
u)
           in Jacobian
qYX Jacobian -> Jacobian -> Jacobian
forall a. Fractional a => a -> a -> a
/ Jacobian
qXY
      j :: Jacobian
j = case Maybe (a -> Double -> Jacobian)
mJac of
        Maybe (a -> Double -> Jacobian)
Nothing -> Jacobian
1.0
        Just a -> Double -> Jacobian
fJac -> a -> Double -> Jacobian
fJac a
x Double
u
  (a, Jacobian, Jacobian) -> IO (a, Jacobian, Jacobian)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
x a -> Double -> a
`f` Double
u, Jacobian
r, Jacobian
j)
{-# INLINEABLE genericContinuous #-}

-- | Generic function to create proposals for discrete parameters (e.g., 'Int').
--
-- See 'genericContinuous'.
genericDiscrete ::
  (DiscreteDistr d, DiscreteGen d) =>
  -- | Probability distribution.
  d ->
  -- | Forward operator.
  --
  -- For example, (+), so that x + dx = x'.
  (a -> Int -> a) ->
  -- | Inverse operator \(g\) of the auxiliary variable.
  --
  -- For example, 'negate', so that x' + (negate dx) = x.
  --
  -- Only required for biased proposals.
  Maybe (Int -> Int) ->
  ProposalSimple a
genericDiscrete :: d -> (a -> Int -> a) -> Maybe (Int -> Int) -> ProposalSimple a
genericDiscrete d
d a -> Int -> a
f Maybe (Int -> Int)
mfInv a
x GenIO
g = do
  Int
u <- d -> Gen RealWorld -> IO Int
forall d g (m :: * -> *).
(DiscreteGen d, StatefulGen g m) =>
d -> g -> m Int
genDiscreteVar d
d Gen RealWorld
GenIO
g
  let r :: Jacobian
r = case Maybe (Int -> Int)
mfInv of
        Maybe (Int -> Int)
Nothing -> Jacobian
1.0
        Just Int -> Int
fInv ->
          let qXY :: Jacobian
qXY = Double -> Jacobian
forall a. a -> Log a
Exp (Double -> Jacobian) -> Double -> Jacobian
forall a b. (a -> b) -> a -> b
$ d -> Int -> Double
forall d. DiscreteDistr d => d -> Int -> Double
logProbability d
d Int
u
              qYX :: Jacobian
qYX = Double -> Jacobian
forall a. a -> Log a
Exp (Double -> Jacobian) -> Double -> Jacobian
forall a b. (a -> b) -> a -> b
$ d -> Int -> Double
forall d. DiscreteDistr d => d -> Int -> Double
logProbability d
d (Int -> Int
fInv Int
u)
           in Jacobian
qYX Jacobian -> Jacobian -> Jacobian
forall a. Fractional a => a -> a -> a
/ Jacobian
qXY
  (a, Jacobian, Jacobian) -> IO (a, Jacobian, Jacobian)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
x a -> Int -> a
`f` Int
u, Jacobian
r, Jacobian
1.0)
{-# INLINEABLE genericDiscrete #-}