{-# LANGUAGE TemplateHaskell #-}

-- |
-- Module      :  Mcmc.Proposal.Simplex
-- Description :  Proposals on simplices
-- Copyright   :  2021 Dominik Schrempf
-- License     :  GPL-3.0-or-later
--
-- Maintainer  :  dominik.schrempf@gmail.com
-- Stability   :  unstable
-- Portability :  portable
--
-- Creation date: Mon Oct 19 15:32:31 2020.
module Mcmc.Proposal.Simplex
  ( -- * Elements of simplices
    Simplex (toVector),
    simplexUniform,
    simplexFromVector,

    -- * Proposals on simplices
    dirichlet,
    beta,
  )
where

import Data.Aeson
import Data.Aeson.TH
import qualified Data.Vector.Unboxed as V
import Mcmc.Proposal
import Mcmc.Statistics.Types
import Numeric.Log
import Statistics.Distribution
import Statistics.Distribution.Beta
import Statistics.Distribution.Dirichlet

-- | An element of a simplex.
--
-- A vector of non-negative values summing to one.
--
-- The nomenclature is not very consistent, because a K-dimensional simplex is
-- usually considered to be the set containing all @K@-dimensional vectors with
-- non-negative elements that sum to 1.0. However, I couldn't come up with a
-- better name. Maybe @SimplexElement@, but that was too long.
newtype Simplex = SimplexUnsafe {Simplex -> Vector TuningParameter
toVector :: V.Vector Double}
  deriving (Simplex -> Simplex -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Simplex -> Simplex -> Bool
$c/= :: Simplex -> Simplex -> Bool
== :: Simplex -> Simplex -> Bool
$c== :: Simplex -> Simplex -> Bool
Eq, Dimension -> Simplex -> ShowS
[Simplex] -> ShowS
Simplex -> String
forall a.
(Dimension -> a -> ShowS)
-> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Simplex] -> ShowS
$cshowList :: [Simplex] -> ShowS
show :: Simplex -> String
$cshow :: Simplex -> String
showsPrec :: Dimension -> Simplex -> ShowS
$cshowsPrec :: Dimension -> Simplex -> ShowS
Show)

$(deriveJSON defaultOptions ''Simplex)

-- Tolerance.
eps :: Double
eps :: TuningParameter
eps = TuningParameter
1e-14

-- Check if vector is normalized with tolerance 'eps'.
isNormalized :: V.Vector Double -> Bool
isNormalized :: Vector TuningParameter -> Bool
isNormalized Vector TuningParameter
v
  | forall a. Num a => a -> a
abs (forall a. (Unbox a, Num a) => Vector a -> a
V.sum Vector TuningParameter
v forall a. Num a => a -> a -> a
- TuningParameter
1.0) forall a. Ord a => a -> a -> Bool
> TuningParameter
eps = Bool
False
  | Bool
otherwise = Bool
True

-- Check if vector contains negative elements.
isNegative :: V.Vector Double -> Bool
isNegative :: Vector TuningParameter -> Bool
isNegative = forall a. Unbox a => (a -> Bool) -> Vector a -> Bool
V.any (forall a. Ord a => a -> a -> Bool
< TuningParameter
0)

-- | Ensure that the value vector is an element of a simplex.
--
-- Return 'Left' if:
-- - The value vector is empty.
-- - The value vector contains negative elements.
-- - The value vector is not normalized.
simplexFromVector :: V.Vector Double -> Either String Simplex
simplexFromVector :: Vector TuningParameter -> Either String Simplex
simplexFromVector Vector TuningParameter
v
  | forall a. Unbox a => Vector a -> Bool
V.null Vector TuningParameter
v = forall a b. a -> Either a b
Left String
"simplexFromVector: Vector is empty."
  | Vector TuningParameter -> Bool
isNegative Vector TuningParameter
v = forall a b. a -> Either a b
Left String
"simplexFromVector: Vector contains negative elements."
  | Bool -> Bool
not (Vector TuningParameter -> Bool
isNormalized Vector TuningParameter
v) = forall a b. a -> Either a b
Left String
"simplexFromVector: Vector is not normalized."
  | Bool
otherwise = forall a b. b -> Either a b
Right forall a b. (a -> b) -> a -> b
$ Vector TuningParameter -> Simplex
SimplexUnsafe Vector TuningParameter
v

-- | Create the uniform element of the K-dimensional simplex.
--
-- Set all values to \(1/D\).
simplexUniform :: Dimension -> Simplex
--                 Should never fail.
simplexUniform :: Dimension -> Simplex
simplexUniform Dimension
k = forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either forall a. HasCallStack => String -> a
error forall a. a -> a
id forall a b. (a -> b) -> a -> b
$ Vector TuningParameter -> Either String Simplex
simplexFromVector forall a b. (a -> b) -> a -> b
$ forall a. Unbox a => Dimension -> a -> Vector a
V.replicate Dimension
k (TuningParameter
1.0 forall a. Fractional a => a -> a -> a
/ forall a b. (Integral a, Num b) => a -> b
fromIntegral Dimension
k)

-- Tuning function is inverted (high alpha means small steps).
getTuningFunction :: TuningParameter -> (TuningParameter -> TuningParameter)
getTuningFunction :: TuningParameter -> TuningParameter -> TuningParameter
getTuningFunction TuningParameter
t = (forall a. Fractional a => a -> a -> a
/ TuningParameter
t'')
  where
    -- Start with small steps.
    t' :: TuningParameter
t' = TuningParameter
t forall a. Fractional a => a -> a -> a
/ TuningParameter
10000
    -- Extremely small tuning parameters lead to numeric overflow. The square
    -- root pulls the tuning parameter closer to 1.0. However, overflow may
    -- still occur (the involved gamma functions grow faster than the
    -- exponential). I did not observe numeric underflow in my tests.
    t'' :: TuningParameter
t'' = forall a. Floating a => a -> a
sqrt TuningParameter
t'

-- The tuning parameter is proportional to the inverted mean of the shape
-- parameter values.
--
-- The values determining the proposal size have been set using an example
-- analysis. They are good values for this analysis, but may fail for other
-- analyses.
dirichletPFunction :: TuningParameter -> PFunction Simplex
dirichletPFunction :: TuningParameter -> PFunction Simplex
dirichletPFunction TuningParameter
t (SimplexUnsafe Vector TuningParameter
xs) IOGenM StdGen
g = do
  -- If @t@ is high and above 1.0, the parameter vector will be low, and the
  -- variance will be high. If @t@ is low and below 1.0, the parameter vector
  -- will be high, and the Dirichlet distribution will be very concentrated with
  -- low variance.
  let ddXs :: DirichletDistribution
ddXs = forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either forall a. HasCallStack => String -> a
error forall a. a -> a
id forall a b. (a -> b) -> a -> b
$ Vector TuningParameter -> Either String DirichletDistribution
dirichletDistribution forall a b. (a -> b) -> a -> b
$ forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
V.map TuningParameter -> TuningParameter
tf Vector TuningParameter
xs
  Vector TuningParameter
ys <- forall g (m :: * -> *).
StatefulGen g m =>
DirichletDistribution -> g -> m (Vector TuningParameter)
dirichletSample DirichletDistribution
ddXs IOGenM StdGen
g
  -- Have to check if parameters are valid (because zeroes do occur).
  let eitherDdYs :: Either String DirichletDistribution
eitherDdYs = Vector TuningParameter -> Either String DirichletDistribution
dirichletDistribution forall a b. (a -> b) -> a -> b
$ forall a b. (Unbox a, Unbox b) => (a -> b) -> Vector a -> Vector b
V.map TuningParameter -> TuningParameter
tf Vector TuningParameter
ys
  let r :: Jacobian
r = case Either String DirichletDistribution
eitherDdYs of
        -- Set ratio to 0; so that the proposal will not be accepted.
        Left String
_ -> Jacobian
0
        Right DirichletDistribution
ddYs -> DirichletDistribution -> Vector TuningParameter -> Jacobian
dirichletDensity DirichletDistribution
ddYs Vector TuningParameter
xs forall a. Fractional a => a -> a -> a
/ DirichletDistribution -> Vector TuningParameter -> Jacobian
dirichletDensity DirichletDistribution
ddXs Vector TuningParameter
ys
  -- I do not think a Jacobian is necessary in this case. I do know that if a
  -- subset of states is updated a Jacobian would be necessary.
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. a -> Jacobian -> Jacobian -> PResult a
Propose (Vector TuningParameter -> Simplex
SimplexUnsafe Vector TuningParameter
ys) Jacobian
r Jacobian
1.0, forall a. Maybe a
Nothing)
  where
    tf :: TuningParameter -> TuningParameter
tf = TuningParameter -> TuningParameter -> TuningParameter
getTuningFunction TuningParameter
t

-- | Dirichlet proposal on a simplex.
--
-- For a given element of a K-dimensional simplex, propose a new element of the
-- K-dimensional simplex. The new element is sampled from the multivariate
-- Dirichlet distribution with parameter vector being proportional to the
-- current element of the simplex.
--
-- The tuning parameter is used to determine the concentration of the Dirichlet
-- distribution: the lower the tuning parameter, the higher the concentration.
--
-- The proposal dimension, which is the dimension of the simplex, is used to
-- determine the optimal acceptance rate.
--
-- For high dimensional simplices, this proposal may have low acceptance rates.
-- In this case, please see the coordinate wise 'beta' proposal.
dirichlet :: PDimension -> PName -> PWeight -> Tune -> Proposal Simplex
dirichlet :: PDimension -> PName -> PWeight -> Tune -> Proposal Simplex
dirichlet = forall a.
PDescription
-> (TuningParameter -> PFunction a)
-> PSpeed
-> PDimension
-> PName
-> PWeight
-> Tune
-> Proposal a
createProposal (String -> PDescription
PDescription String
"Dirichlet") TuningParameter -> PFunction Simplex
dirichletPFunction PSpeed
PFast

-- The tuning parameter is the inverted mean of the shape values.
--
-- The values determining the proposal size have been set using an example
-- analysis. They are good values for this analysis, but may fail for other
-- analyses.
--
-- See also the 'dirichlet' proposal.
betaPFunction :: Dimension -> TuningParameter -> PFunction Simplex
betaPFunction :: Dimension -> TuningParameter -> PFunction Simplex
betaPFunction Dimension
i TuningParameter
t (SimplexUnsafe Vector TuningParameter
xs) IOGenM StdGen
g = do
  -- Shape parameters of beta distribution. Do not assume that the sum of the
  -- elements of 'xs' is 1.0, because then repeated proposals let the sum of the
  -- vector diverge.
  let aX :: TuningParameter
aX = TuningParameter
xI
      bX :: TuningParameter
bX = TuningParameter
xsSum forall a. Num a => a -> a -> a
- TuningParameter
xI
      bdXI :: BetaDistribution
bdXI = TuningParameter -> TuningParameter -> BetaDistribution
betaDistr (TuningParameter -> TuningParameter
tf TuningParameter
aX) (TuningParameter -> TuningParameter
tf TuningParameter
bX)
  -- New value of element i.
  TuningParameter
yI <- forall d g (m :: * -> *).
(ContGen d, StatefulGen g m) =>
d -> g -> m TuningParameter
genContVar BetaDistribution
bdXI IOGenM StdGen
g
  -- Shape parameters of beta distribution.
  let aY :: TuningParameter
aY = TuningParameter
yI
      bY :: TuningParameter
bY = TuningParameter
1.0 forall a. Num a => a -> a -> a
- TuningParameter
yI
      eitherBdYI :: Maybe BetaDistribution
eitherBdYI = TuningParameter -> TuningParameter -> Maybe BetaDistribution
betaDistrE (TuningParameter -> TuningParameter
tf TuningParameter
aY) (TuningParameter -> TuningParameter
tf TuningParameter
bY)
  -- See 'dirichlet', which has the same construct.
  let r :: Jacobian
r = case Maybe BetaDistribution
eitherBdYI of
        Maybe BetaDistribution
Nothing -> Jacobian
0
        Just BetaDistribution
bdYI -> forall a. a -> Log a
Exp forall a b. (a -> b) -> a -> b
$ forall d. ContDistr d => d -> TuningParameter -> TuningParameter
logDensity BetaDistribution
bdYI TuningParameter
xI forall a. Num a => a -> a -> a
- forall d. ContDistr d => d -> TuningParameter -> TuningParameter
logDensity BetaDistribution
bdXI TuningParameter
yI
      -- The absolute value of the determinant of the Jacobian. Derivation takes
      -- a while...
      ja1 :: TuningParameter
ja1 = TuningParameter
bY forall a. Fractional a => a -> a -> a
/ TuningParameter
bX
      jac :: Jacobian
jac = forall a. a -> Log a
Exp forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Unbox a => Vector a -> Dimension
V.length Vector TuningParameter
xs forall a. Num a => a -> a -> a
- Dimension
2) forall a. Num a => a -> a -> a
* forall a. Floating a => a -> a
log TuningParameter
ja1
  -- Construct new vector.
  let -- Normalization function for other elements.
      -- nf x = x * bY / bX
      --
      -- It turns out, that this factor is also needed to compute the determinant
      -- of the Jacobian above.
      nf :: TuningParameter -> TuningParameter
nf TuningParameter
x = TuningParameter
x forall a. Num a => a -> a -> a
* TuningParameter
ja1
      ys :: Vector TuningParameter
ys = forall a. Unbox a => Dimension -> (Dimension -> a) -> Vector a
V.generate (forall a. Unbox a => Vector a -> Dimension
V.length Vector TuningParameter
xs) (\Dimension
j -> if Dimension
i forall a. Eq a => a -> a -> Bool
== Dimension
j then TuningParameter
yI else TuningParameter -> TuningParameter
nf (Vector TuningParameter
xs forall a. Unbox a => Vector a -> Dimension -> a
V.! Dimension
j))
      y :: Simplex
y = forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either forall a. HasCallStack => String -> a
error forall a. a -> a
id forall a b. (a -> b) -> a -> b
$ Vector TuningParameter -> Either String Simplex
simplexFromVector Vector TuningParameter
ys
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. a -> Jacobian -> Jacobian -> PResult a
Propose Simplex
y Jacobian
r Jacobian
jac, forall a. Maybe a
Nothing)
  where
    xI :: TuningParameter
xI = Vector TuningParameter
xs forall a. Unbox a => Vector a -> Dimension -> a
V.! Dimension
i
    xsSum :: TuningParameter
xsSum = forall a. (Unbox a, Num a) => Vector a -> a
V.sum Vector TuningParameter
xs
    tf :: TuningParameter -> TuningParameter
tf = TuningParameter -> TuningParameter -> TuningParameter
getTuningFunction TuningParameter
t

-- | Beta proposal on a specific coordinate @i@ on a simplex.
--
-- For a given element of a K-dimensional simplex, propose a new element of the
-- K-dimensional simplex. The coordinate @i@ of the new element is sampled from
-- the beta distribution. The other coordinates are normalized such that the
-- values sum to 1.0. The parameters of the beta distribution are chosen such
-- that the expected value of the beta distribution is the value of the old
-- coordinate.
--
-- The tuning parameter is used to determine the concentration of the beta
-- distribution: the lower the tuning parameter, the higher the concentration.
--
-- See also the 'dirichlet' proposal.
--
-- No "out of bounds" checks are performed during compile time. Run time errors
-- can occur if @i@ is negative, or if @i-1@ is larger than the length of the
-- element vector of the simplex.
--
-- This proposal has been assigned a dimension of 2. See the discussion at
-- 'PDimension'.
beta :: Dimension -> PName -> PWeight -> Tune -> Proposal Simplex
beta :: Dimension -> PName -> PWeight -> Tune -> Proposal Simplex
beta Dimension
i = forall a.
PDescription
-> (TuningParameter -> PFunction a)
-> PSpeed
-> PDimension
-> PName
-> PWeight
-> Tune
-> Proposal a
createProposal PDescription
description (Dimension -> TuningParameter -> PFunction Simplex
betaPFunction Dimension
i) PSpeed
PFast (Dimension -> PDimension
PDimension Dimension
2)
  where
    description :: PDescription
description = String -> PDescription
PDescription forall a b. (a -> b) -> a -> b
$ String
"Beta; coordinate: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Dimension
i