{-# LANGUAGE TupleSections #-}
{-# LANGUAGE UndecidableInstances #-}

-- |
-- Module      :  Mcmc.Proposal.Hamiltonian
-- Description :  Hamiltonian Monte Carlo proposal
-- Copyright   :  (c) 2021 Dominik Schrempf
-- License     :  GPL-3.0-or-later
--
-- Maintainer  :  dominik.schrempf@gmail.com
-- Stability   :  experimental
-- Portability :  portable
--
-- Creation date: Mon Jul  5 12:59:42 2021.
--
-- The Hamiltonian Monte Carlo (HMC) proposal.
--
-- For references, see:
--
-- - [1] Chapter 5 of Handbook of Monte Carlo: Neal, R. M., MCMC Using
--   Hamiltonian Dynamics, In S. Brooks, A. Gelman, G. Jones, & X. Meng (Eds.),
--   Handbook of Markov Chain Monte Carlo (2011), CRC press.
--
-- - [2] Gelman, A., Carlin, J. B., Stern, H. S., & Rubin, D. B., Bayesian data
--   analysis (2014), CRC Press.
--
-- - [3] Review by Betancourt and notes: Betancourt, M., A conceptual
--   introduction to Hamiltonian Monte Carlo, arXiv, 1701–02434 (2017).
--
-- NOTE on implementation:
--
-- - The implementation assumes the existence of the 'Gradient'. Like so, the
--   user can use automatic or manual differentiation, depending on the problem
--   at hand.
--
-- - The Hamiltonian proposal acts on a vector of storable 'Values'. Functions
--   converting the state to and from this vector have to be provided. See
--   'HSettings'.
--
-- - The desired acceptance rate is 0.65, although the dimension of the proposal
--   is high.
--
-- - The speed of this proposal can change drastically when tuned because the
--   leapfrog trajectory length is changed.
--
-- - The Hamiltonian proposal is agnostic of the actual prior and likelihood
--   functions, and so, points with zero posterior probability cannot be
--   detected. This affects models with constrained parameters. See Gelman p.
--   303. This problem can be ameliorated by providing a 'Validate' function so
--   that the proposal can gracefully fail as soon as the state becomes invalid.
module Mcmc.Proposal.Hamiltonian
  ( Values,
    Gradient,
    Validate,
    Masses,
    LeapfrogTrajectoryLength,
    LeapfrogScalingFactor,
    HTuneLeapfrog (..),
    HTuneMasses (..),
    HTune (..),
    HSettings (..),
    hamiltonian,
  )
where

import qualified Data.Vector as VB
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Unboxed as VU
import Mcmc.Proposal
import qualified Numeric.LinearAlgebra as L
import Numeric.Log
import Numeric.MathFunctions.Constants
import qualified Statistics.Covariance as S
import qualified Statistics.Function as S
import qualified Statistics.Sample as S
import System.Random.MWC

-- TODO: No-U-turn sampler (NUTS). Ameliorates necessity to determine the
-- leapfrog trajectory length L. (I think this is a necessary extension.)

-- TODO: Riemannian adaptation: State-dependent mass matrix. (Seems a little bit
-- of an overkill.)

-- | The Hamiltonian proposal acts on a vector of floating point values.
type Values = L.Vector Double

-- | Gradient of the log posterior function.
--
-- The gradient has to be provided for the complete state. The reason is that
-- the gradient may change if parameters untouched by the Hamiltonian proposal
-- are altered by other proposals.
type Gradient a = a -> a

-- | Function validating the state.
--
-- Useful if parameters are constrained.
--
-- Also the validity of the state may depend on parameters untouched by the
-- Hamiltonian proposal.
type Validate a = a -> Bool

-- | Parameter mass matrix.
--
-- The masses roughly describe how reluctant the particles move through the
-- state space. If a parameter has higher mass, the momentum in this direction
-- will be changed less by the provided gradient, than when the same parameter
-- has lower mass. Off-diagonal entries describe the covariance structure. If
-- two parameters are negatively correlated, their generated initial momenta are
-- likely to have opposite signs.
--
-- The proposal is more efficient if masses are assigned according to the
-- inverse (co)-variance structure of the posterior function. That is,
-- parameters changing on larger scales should have lower masses than parameters
-- changing on lower scales. In particular, the optimal entries of the diagonal
-- of the mass matrix are the inverted variances of the parameters distributed
-- according to the posterior function.
--
-- Of course, the scales of the parameters of the posterior function are usually
-- unknown. Often, it is sufficient to
--
-- - set the diagonal entries of the mass matrix to identical values roughly
--   scaled with the inverted estimated average variance of the posterior
--   function; or even to
--
-- - set all diagonal entries of the mass matrix to 1.0, and all other entries
--   to 0.0, and trust the tuning algorithm (see 'HTune') to find the correct
--   values.
type Masses = L.Herm Double

-- | Mean leapfrog trajectory length \(L\).
--
-- Number of leapfrog steps per proposal.
--
-- To avoid problems with ergodicity, the actual number of leapfrog steps is
-- sampled per proposal from a discrete uniform distribution over the interval
-- \([\text{floor}(0.8L),\text{ceiling}(1.2L)]\).
--
-- For a discussion of ergodicity and reasons why randomization is important,
-- see [1] p. 15; also mentioned in [2] p. 304.
--
-- Usually set to 10, but larger values may be desirable.
--
-- NOTE: To avoid errors, the left bound has an additional hard minimum of 1,
-- and the right bound is required to be larger equal than the left bound.
--
-- NOTE: Call 'error' if value is less than 1.
type LeapfrogTrajectoryLength = Int

-- | Mean of leapfrog scaling factor \(\epsilon\).
--
-- Determines the size of each leapfrog step.
--
-- To avoid problems with ergodicity, the actual leapfrog scaling factor is
-- sampled per proposal from a continuous uniform distribution over the interval
-- \((0.8\epsilon,1.2\epsilon]\).
--
-- For a discussion of ergodicity and reasons why randomization is important,
-- see [1] p. 15; also mentioned in [2] p. 304.
--
-- Usually set such that \( L \epsilon = 1.0 \), but smaller values may be
-- required if acceptance rates are low.
--
-- NOTE: Call 'error' if value is zero or negative.
type LeapfrogScalingFactor = Double

-- Internal. Values; target state containing parameters.
type Positions = Values

-- Internal. Momenta of the parameters.
type Momenta = L.Vector Double

-- | Tune leapfrog parameters?
data HTuneLeapfrog
  = HNoTuneLeapfrog
  | -- | We expect that the larger the leapfrog scaling factor the lower the
    -- acceptance ratio. Consequently, if the acceptance rate is too low, the
    -- leapfrog scaling factor is decreased and vice versa. Further, the leapfrog
    -- trajectory length is scaled such that the product of the leapfrog scaling
    -- factor and leapfrog trajectory length stays roughly constant.
    HTuneLeapfrog
  deriving (HTuneLeapfrog -> HTuneLeapfrog -> Bool
(HTuneLeapfrog -> HTuneLeapfrog -> Bool)
-> (HTuneLeapfrog -> HTuneLeapfrog -> Bool) -> Eq HTuneLeapfrog
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: HTuneLeapfrog -> HTuneLeapfrog -> Bool
$c/= :: HTuneLeapfrog -> HTuneLeapfrog -> Bool
== :: HTuneLeapfrog -> HTuneLeapfrog -> Bool
$c== :: HTuneLeapfrog -> HTuneLeapfrog -> Bool
Eq, Int -> HTuneLeapfrog -> ShowS
[HTuneLeapfrog] -> ShowS
HTuneLeapfrog -> String
(Int -> HTuneLeapfrog -> ShowS)
-> (HTuneLeapfrog -> String)
-> ([HTuneLeapfrog] -> ShowS)
-> Show HTuneLeapfrog
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [HTuneLeapfrog] -> ShowS
$cshowList :: [HTuneLeapfrog] -> ShowS
show :: HTuneLeapfrog -> String
$cshow :: HTuneLeapfrog -> String
showsPrec :: Int -> HTuneLeapfrog -> ShowS
$cshowsPrec :: Int -> HTuneLeapfrog -> ShowS
Show)

-- | Tune masses?
--
-- The masses are tuned according to the (co)variances of the parameters
-- obtained from the posterior distribution over the last auto tuning interval.
data HTuneMasses
  = HNoTuneMasses
  | -- | Diagonal only: The variances of the parameters are calculated and the
    -- masses are amended using the old masses and the inverted variances. If, for
    -- a specific coordinate, the sample size is 60 or lower, or if the calculated
    -- variance is out of predefined bounds [1e-6, 1e6], the mass of the affected
    -- position is not changed.
    HTuneDiagonalMassesOnly
  | -- | All masses: The covariance matrix of the parameters is estimated and the
    -- inverted matrix (sometimes called precision matrix) is used as mass matrix.
    -- This procedure is error prone, but models with high correlations between
    -- parameters it is necessary to tune off-diagonal entries. The full mass
    -- matrix is only tuned if more than 200 samples are available. For these
    -- reasons, when tuning all masses it is recommended to use tuning settings
    -- such as
    --
    -- @
    -- BurnInWithCustomAutoTuning ([10, 20 .. 200] ++ replicate 5 500)
    -- @
    HTuneAllMasses
  deriving (HTuneMasses -> HTuneMasses -> Bool
(HTuneMasses -> HTuneMasses -> Bool)
-> (HTuneMasses -> HTuneMasses -> Bool) -> Eq HTuneMasses
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: HTuneMasses -> HTuneMasses -> Bool
$c/= :: HTuneMasses -> HTuneMasses -> Bool
== :: HTuneMasses -> HTuneMasses -> Bool
$c== :: HTuneMasses -> HTuneMasses -> Bool
Eq, Int -> HTuneMasses -> ShowS
[HTuneMasses] -> ShowS
HTuneMasses -> String
(Int -> HTuneMasses -> ShowS)
-> (HTuneMasses -> String)
-> ([HTuneMasses] -> ShowS)
-> Show HTuneMasses
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [HTuneMasses] -> ShowS
$cshowList :: [HTuneMasses] -> ShowS
show :: HTuneMasses -> String
$cshow :: HTuneMasses -> String
showsPrec :: Int -> HTuneMasses -> ShowS
$cshowsPrec :: Int -> HTuneMasses -> ShowS
Show)

-- | Tuning settings.
data HTune = HTune HTuneLeapfrog HTuneMasses
  deriving (HTune -> HTune -> Bool
(HTune -> HTune -> Bool) -> (HTune -> HTune -> Bool) -> Eq HTune
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: HTune -> HTune -> Bool
$c/= :: HTune -> HTune -> Bool
== :: HTune -> HTune -> Bool
$c== :: HTune -> HTune -> Bool
Eq, Int -> HTune -> ShowS
[HTune] -> ShowS
HTune -> String
(Int -> HTune -> ShowS)
-> (HTune -> String) -> ([HTune] -> ShowS) -> Show HTune
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [HTune] -> ShowS
$cshowList :: [HTune] -> ShowS
show :: HTune -> String
$cshow :: HTune -> String
showsPrec :: Int -> HTune -> ShowS
$cshowsPrec :: Int -> HTune -> ShowS
Show)

-- | Specifications of the Hamilton Monte Carlo proposal.
data HSettings a = HSettings
  { -- | Extract values to be manipulated by the Hamiltonian proposal from the
    -- state.
    HSettings a -> a -> Values
hToVector :: a -> Values,
    -- | Put those values back into the state.
    HSettings a -> a -> Values -> a
hFromVectorWith :: a -> Values -> a,
    HSettings a -> Gradient a
hGradient :: Gradient a,
    HSettings a -> Maybe (Validate a)
hMaybeValidate :: Maybe (Validate a),
    HSettings a -> Masses
hMasses :: Masses,
    HSettings a -> Int
hLeapfrogTrajectoryLength :: LeapfrogTrajectoryLength,
    HSettings a -> LeapfrogScalingFactor
hLeapfrogScalingFactor :: LeapfrogScalingFactor,
    HSettings a -> HTune
hTune :: HTune
  }

checkHSettings :: Eq a => a -> HSettings a -> Maybe String
checkHSettings :: a -> HSettings a -> Maybe String
checkHSettings a
x (HSettings a -> Values
toVec a -> Values -> a
fromVec Gradient a
_ Maybe (Validate a)
_ Masses
masses Int
l LeapfrogScalingFactor
eps HTune
_)
  | (LeapfrogScalingFactor -> Bool) -> [LeapfrogScalingFactor] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (LeapfrogScalingFactor -> LeapfrogScalingFactor -> Bool
forall a. Ord a => a -> a -> Bool
<= LeapfrogScalingFactor
0) [LeapfrogScalingFactor]
diagonalMasses = String -> Maybe String
forall a. a -> Maybe a
Just String
"checkHSettings: Some diagonal entries of the mass matrix are zero or negative."
  | Int
nrows Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
ncols = String -> Maybe String
forall a. a -> Maybe a
Just String
"checkHSettings: Mass matrix is not square."
  | a -> Values -> a
fromVec a
x Values
xVec a -> Validate a
forall a. Eq a => a -> a -> Bool
/= a
x = String -> Maybe String
forall a. a -> Maybe a
Just String
"checkHSettings: 'fromVectorWith x (toVector x) /= x' for sample state."
  | Values -> IndexOf Vector
forall (c :: * -> *) t. Container c t => c t -> IndexOf c
L.size Values
xVec Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
nrows = String -> Maybe String
forall a. a -> Maybe a
Just String
"checkHSettings: Mass matrix has different size than 'toVector x', where x is sample state."
  | Int
l Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1 = String -> Maybe String
forall a. a -> Maybe a
Just String
"checkHSettings: Leapfrog trajectory length is zero or negative."
  | LeapfrogScalingFactor
eps LeapfrogScalingFactor -> LeapfrogScalingFactor -> Bool
forall a. Ord a => a -> a -> Bool
<= LeapfrogScalingFactor
0 = String -> Maybe String
forall a. a -> Maybe a
Just String
"checkHSettings: Leapfrog scaling factor is zero or negative."
  | Bool
otherwise = Maybe String
forall a. Maybe a
Nothing
  where
    ms :: Matrix LeapfrogScalingFactor
ms = Masses -> Matrix LeapfrogScalingFactor
forall t. Herm t -> Matrix t
L.unSym Masses
masses
    diagonalMasses :: [LeapfrogScalingFactor]
diagonalMasses = Values -> [LeapfrogScalingFactor]
forall a. Storable a => Vector a -> [a]
L.toList (Values -> [LeapfrogScalingFactor])
-> Values -> [LeapfrogScalingFactor]
forall a b. (a -> b) -> a -> b
$ Matrix LeapfrogScalingFactor -> Values
forall t. Element t => Matrix t -> Vector t
L.takeDiag Matrix LeapfrogScalingFactor
ms
    nrows :: Int
nrows = Matrix LeapfrogScalingFactor -> Int
forall t. Matrix t -> Int
L.rows Matrix LeapfrogScalingFactor
ms
    ncols :: Int
ncols = Matrix LeapfrogScalingFactor -> Int
forall t. Matrix t -> Int
L.cols Matrix LeapfrogScalingFactor
ms
    xVec :: Values
xVec = a -> Values
toVec a
x

-- Internal. Mean vector containing zeroes.
type HMu = L.Vector Double

-- Internal. Symmetric, inverted mass matrix.
type HMassesInv = L.Herm Double

-- Internal. Symmetric, inverted mass matrix scaled with the leapfrog step size
-- epsilon.
type HMassesInvEps = L.Herm Double

-- Internal. Logarithm of the determinant of the mass matrix.
type HLogDetMasses = Double

-- Internal data type containing memoized values.
data HData = HData
  { HData -> Values
_hMu :: HMu,
    HData -> Masses
_hMassesInv :: HMassesInv,
    HData -> Masses
_hMassesInvEps :: HMassesInvEps,
    HData -> LeapfrogScalingFactor
_hLogDetMasses :: HLogDetMasses
  }

-- Call 'error' if the determinant of the covariance matrix is negative.
getHData :: HSettings a -> HData
getHData :: HSettings a -> HData
getHData HSettings a
s =
  -- The multivariate normal distribution requires a positive definite matrix
  -- with positive determinant.
  if LeapfrogScalingFactor
sign LeapfrogScalingFactor -> LeapfrogScalingFactor -> Bool
forall a. Eq a => a -> a -> Bool
== LeapfrogScalingFactor
1.0
    then Values -> Masses -> Masses -> LeapfrogScalingFactor -> HData
HData Values
mu Masses
massesInvH Masses
massesInvEpsH LeapfrogScalingFactor
logDetMasses
    else
      let msg :: String
msg =
            String
"hamiltonianSimple: Determinant of covariance matrix is negative."
              String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
" The logarithm of the absolute value of the determinant is: "
              String -> ShowS
forall a. Semigroup a => a -> a -> a
<> LeapfrogScalingFactor -> String
forall a. Show a => a -> String
show LeapfrogScalingFactor
logDetMasses
              String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
"."
       in String -> HData
forall a. HasCallStack => String -> a
error String
msg
  where
    ms :: Masses
ms = HSettings a -> Masses
forall a. HSettings a -> Masses
hMasses HSettings a
s
    nrows :: Int
nrows = Matrix LeapfrogScalingFactor -> Int
forall t. Matrix t -> Int
L.rows (Matrix LeapfrogScalingFactor -> Int)
-> Matrix LeapfrogScalingFactor -> Int
forall a b. (a -> b) -> a -> b
$ Masses -> Matrix LeapfrogScalingFactor
forall t. Herm t -> Matrix t
L.unSym Masses
ms
    mu :: Values
mu = [LeapfrogScalingFactor] -> Values
forall a. Storable a => [a] -> Vector a
L.fromList ([LeapfrogScalingFactor] -> Values)
-> [LeapfrogScalingFactor] -> Values
forall a b. (a -> b) -> a -> b
$ Int -> LeapfrogScalingFactor -> [LeapfrogScalingFactor]
forall a. Int -> a -> [a]
replicate Int
nrows LeapfrogScalingFactor
0.0
    (Matrix LeapfrogScalingFactor
massesInv, (LeapfrogScalingFactor
logDetMasses, LeapfrogScalingFactor
sign)) = Matrix LeapfrogScalingFactor
-> (Matrix LeapfrogScalingFactor,
    (LeapfrogScalingFactor, LeapfrogScalingFactor))
forall t. Field t => Matrix t -> (Matrix t, (t, t))
L.invlndet (Matrix LeapfrogScalingFactor
 -> (Matrix LeapfrogScalingFactor,
     (LeapfrogScalingFactor, LeapfrogScalingFactor)))
-> Matrix LeapfrogScalingFactor
-> (Matrix LeapfrogScalingFactor,
    (LeapfrogScalingFactor, LeapfrogScalingFactor))
forall a b. (a -> b) -> a -> b
$ Masses -> Matrix LeapfrogScalingFactor
forall t. Herm t -> Matrix t
L.unSym Masses
ms
    -- In theory we can trust that the matrix is symmetric here, because the
    -- inverse of a symmetric matrix is symmetric. However, one may want to
    -- implement a check anyways.
    massesInvH :: Masses
massesInvH = Matrix LeapfrogScalingFactor -> Masses
forall t. Matrix t -> Herm t
L.trustSym Matrix LeapfrogScalingFactor
massesInv
    eps :: LeapfrogScalingFactor
eps = HSettings a -> LeapfrogScalingFactor
forall a. HSettings a -> LeapfrogScalingFactor
hLeapfrogScalingFactor HSettings a
s
    massesInvEpsH :: Masses
massesInvEpsH = LeapfrogScalingFactor -> Masses -> Masses
forall t (c :: * -> *). Linear t c => t -> c t -> c t
L.scale LeapfrogScalingFactor
eps Masses
massesInvH

generateMomenta ::
  -- Provided so that it does not have to be recreated.
  HMu ->
  Masses ->
  GenIO ->
  IO Momenta
generateMomenta :: Values -> Masses -> GenIO -> IO Values
generateMomenta Values
mu Masses
masses GenIO
gen = do
  Int
seed <- Gen RealWorld -> IO Int
forall a g (m :: * -> *). (Uniform a, StatefulGen g m) => g -> m a
uniformM Gen RealWorld
GenIO
gen :: IO Int
  let momenta :: Matrix LeapfrogScalingFactor
momenta = Int -> Int -> Values -> Masses -> Matrix LeapfrogScalingFactor
L.gaussianSample Int
seed Int
1 Values
mu Masses
masses
  Values -> IO Values
forall (m :: * -> *) a. Monad m => a -> m a
return (Values -> IO Values) -> Values -> IO Values
forall a b. (a -> b) -> a -> b
$ Matrix LeapfrogScalingFactor -> Values
forall t. Element t => Matrix t -> Vector t
L.flatten Matrix LeapfrogScalingFactor
momenta

-- Prior distribution of momenta.
--
-- Log of density of multivariate normal distribution with given parameters.
-- https://en.wikipedia.org/wiki/Multivariate_normal_distribution#Density_function.
logDensityMultivariateNormal ::
  -- Mean vector.
  L.Vector Double ->
  -- Inverted covariance matrix.
  L.Herm Double ->
  -- Logarithm of the determinant of the covariance matrix.
  Double ->
  -- Value vector.
  L.Vector Double ->
  Log Double
logDensityMultivariateNormal :: Values
-> Masses
-> LeapfrogScalingFactor
-> Values
-> Log LeapfrogScalingFactor
logDensityMultivariateNormal Values
mu Masses
sigmaInvH LeapfrogScalingFactor
logDetSigma Values
xs =
  LeapfrogScalingFactor -> Log LeapfrogScalingFactor
forall a. a -> Log a
Exp (LeapfrogScalingFactor -> Log LeapfrogScalingFactor)
-> LeapfrogScalingFactor -> Log LeapfrogScalingFactor
forall a b. (a -> b) -> a -> b
$ LeapfrogScalingFactor
c LeapfrogScalingFactor
-> LeapfrogScalingFactor -> LeapfrogScalingFactor
forall a. Num a => a -> a -> a
+ (-LeapfrogScalingFactor
0.5) LeapfrogScalingFactor
-> LeapfrogScalingFactor -> LeapfrogScalingFactor
forall a. Num a => a -> a -> a
* (LeapfrogScalingFactor
logDetSigma LeapfrogScalingFactor
-> LeapfrogScalingFactor -> LeapfrogScalingFactor
forall a. Num a => a -> a -> a
+ ((Values
dxs Values -> Matrix LeapfrogScalingFactor -> Values
forall t. Numeric t => Vector t -> Matrix t -> Vector t
L.<# Matrix LeapfrogScalingFactor
sigmaInv) Values -> Values -> LeapfrogScalingFactor
forall t. Numeric t => Vector t -> Vector t -> t
L.<.> Values
dxs))
  where
    dxs :: Values
dxs = Values
xs Values -> Values -> Values
forall a. Num a => a -> a -> a
- Values
mu
    k :: LeapfrogScalingFactor
k = Int -> LeapfrogScalingFactor
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> LeapfrogScalingFactor) -> Int -> LeapfrogScalingFactor
forall a b. (a -> b) -> a -> b
$ Values -> IndexOf Vector
forall (c :: * -> *) t. Container c t => c t -> IndexOf c
L.size Values
mu
    c :: LeapfrogScalingFactor
c = LeapfrogScalingFactor -> LeapfrogScalingFactor
forall a. Num a => a -> a
negate (LeapfrogScalingFactor -> LeapfrogScalingFactor)
-> LeapfrogScalingFactor -> LeapfrogScalingFactor
forall a b. (a -> b) -> a -> b
$ LeapfrogScalingFactor
m_ln_sqrt_2_pi LeapfrogScalingFactor
-> LeapfrogScalingFactor -> LeapfrogScalingFactor
forall a. Num a => a -> a -> a
* LeapfrogScalingFactor
k
    sigmaInv :: Matrix LeapfrogScalingFactor
sigmaInv = Masses -> Matrix LeapfrogScalingFactor
forall t. Herm t -> Matrix t
L.unSym Masses
sigmaInvH

leapfrog ::
  Gradient Positions ->
  Maybe (Validate Positions) ->
  HMassesInvEps ->
  LeapfrogTrajectoryLength ->
  LeapfrogScalingFactor ->
  Positions ->
  Momenta ->
  -- Maybe (Positions', Momenta'); fail if state is not valid.
  Maybe (Positions, Momenta)
leapfrog :: (Values -> Values)
-> Maybe (Validate Values)
-> Masses
-> Int
-> LeapfrogScalingFactor
-> Values
-> Values
-> Maybe (Values, Values)
leapfrog Values -> Values
grad Maybe (Validate Values)
mVal Masses
hMassesInvEps Int
l LeapfrogScalingFactor
eps Values
theta Values
phi = do
  let -- The first half step of the momenta.
      phiHalf :: Values
phiHalf = LeapfrogScalingFactor
-> (Values -> Values) -> Values -> Values -> Values
leapfrogStepMomenta (LeapfrogScalingFactor
0.5 LeapfrogScalingFactor
-> LeapfrogScalingFactor -> LeapfrogScalingFactor
forall a. Num a => a -> a -> a
* LeapfrogScalingFactor
eps) Values -> Values
grad Values
theta Values
phi
  -- L-1 full steps. This gives the positions theta_{L-1}, and the momenta
  -- phi_{L-1/2}.
  (Values
thetaLM1, Values
phiLM1Half) <- Int -> Maybe (Values, Values) -> Maybe (Values, Values)
forall t.
(Eq t, Num t) =>
t -> Maybe (Values, Values) -> Maybe (Values, Values)
go (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) ((Values, Values) -> Maybe (Values, Values)
forall a. a -> Maybe a
Just (Values
theta, Values
phiHalf))
  -- The last full step of the positions.
  Values
thetaL <- Values -> Maybe Values
valF (Values -> Maybe Values) -> Values -> Maybe Values
forall a b. (a -> b) -> a -> b
$ Masses -> Values -> Values -> Values
leapfrogStepPositions Masses
hMassesInvEps Values
thetaLM1 Values
phiLM1Half
  let -- The last half step of the momenta.
      phiL :: Values
phiL = LeapfrogScalingFactor
-> (Values -> Values) -> Values -> Values -> Values
leapfrogStepMomenta (LeapfrogScalingFactor
0.5 LeapfrogScalingFactor
-> LeapfrogScalingFactor -> LeapfrogScalingFactor
forall a. Num a => a -> a -> a
* LeapfrogScalingFactor
eps) Values -> Values
grad Values
thetaL Values
phiLM1Half
  (Values, Values) -> Maybe (Values, Values)
forall (m :: * -> *) a. Monad m => a -> m a
return (Values
thetaL, Values
phiL)
  where
    valF :: Values -> Maybe Values
valF Values
x = case Maybe (Validate Values)
mVal of
      Maybe (Validate Values)
Nothing -> Values -> Maybe Values
forall a. a -> Maybe a
Just Values
x
      Just Validate Values
f -> if Validate Values
f Values
x then Values -> Maybe Values
forall a. a -> Maybe a
Just Values
x else Maybe Values
forall a. Maybe a
Nothing
    go :: t -> Maybe (Values, Values) -> Maybe (Values, Values)
go t
_ Maybe (Values, Values)
Nothing = Maybe (Values, Values)
forall a. Maybe a
Nothing
    go t
0 (Just (Values
t, Values
p)) = (Values, Values) -> Maybe (Values, Values)
forall a. a -> Maybe a
Just (Values
t, Values
p)
    go t
n (Just (Values
t, Values
p)) =
      let t' :: Values
t' = Masses -> Values -> Values -> Values
leapfrogStepPositions Masses
hMassesInvEps Values
t Values
p
          p' :: Values
p' = LeapfrogScalingFactor
-> (Values -> Values) -> Values -> Values -> Values
leapfrogStepMomenta LeapfrogScalingFactor
eps Values -> Values
grad Values
t' Values
p
          r :: Maybe (Values, Values)
r = (,Values
p') (Values -> (Values, Values))
-> Maybe Values -> Maybe (Values, Values)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Values -> Maybe Values
valF Values
t'
       in t -> Maybe (Values, Values) -> Maybe (Values, Values)
go (t
n t -> t -> t
forall a. Num a => a -> a -> a
- t
1) Maybe (Values, Values)
r

leapfrogStepMomenta ::
  -- Step size (usually a multiple of the leapfrog scaling factor).
  Double ->
  Gradient Positions ->
  -- Current positions.
  Positions ->
  -- Current momenta.
  Momenta ->
  -- New momenta.
  Momenta
leapfrogStepMomenta :: LeapfrogScalingFactor
-> (Values -> Values) -> Values -> Values -> Values
leapfrogStepMomenta LeapfrogScalingFactor
eps Values -> Values
grad Values
theta Values
phi = Values
phi Values -> Values -> Values
forall a. Num a => a -> a -> a
+ LeapfrogScalingFactor -> Values -> Values
forall t (c :: * -> *). Linear t c => t -> c t -> c t
L.scale LeapfrogScalingFactor
eps (Values -> Values
grad Values
theta)

leapfrogStepPositions ::
  HMassesInvEps ->
  -- Current positions.
  Positions ->
  -- Current momenta.
  Momenta ->
  -- New positions.
  Positions
leapfrogStepPositions :: Masses -> Values -> Values -> Values
leapfrogStepPositions Masses
hMassesInvEps Values
theta Values
phi = Values
theta Values -> Values -> Values
forall a. Num a => a -> a -> a
+ (Masses -> Matrix LeapfrogScalingFactor
forall t. Herm t -> Matrix t
L.unSym Masses
hMassesInvEps Matrix LeapfrogScalingFactor -> Values -> Values
forall t. Numeric t => Matrix t -> Vector t -> Vector t
L.#> Values
phi)

massesToTuningParameters :: Masses -> AuxiliaryTuningParameters
massesToTuningParameters :: Masses -> AuxiliaryTuningParameters
massesToTuningParameters = Values -> AuxiliaryTuningParameters
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VB.convert (Values -> AuxiliaryTuningParameters)
-> (Masses -> Values) -> Masses -> AuxiliaryTuningParameters
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Matrix LeapfrogScalingFactor -> Values
forall t. Element t => Matrix t -> Vector t
L.flatten (Matrix LeapfrogScalingFactor -> Values)
-> (Masses -> Matrix LeapfrogScalingFactor) -> Masses -> Values
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Masses -> Matrix LeapfrogScalingFactor
forall t. Herm t -> Matrix t
L.unSym

tuningParametersToMasses ::
  -- Dimension of the mass matrix.
  Int ->
  AuxiliaryTuningParameters ->
  Masses
tuningParametersToMasses :: Int -> AuxiliaryTuningParameters -> Masses
tuningParametersToMasses Int
d = Matrix LeapfrogScalingFactor -> Masses
forall t. Matrix t -> Herm t
L.trustSym (Matrix LeapfrogScalingFactor -> Masses)
-> (AuxiliaryTuningParameters -> Matrix LeapfrogScalingFactor)
-> AuxiliaryTuningParameters
-> Masses
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Values -> Matrix LeapfrogScalingFactor
forall t. Storable t => Int -> Vector t -> Matrix t
L.reshape Int
d (Values -> Matrix LeapfrogScalingFactor)
-> (AuxiliaryTuningParameters -> Values)
-> AuxiliaryTuningParameters
-> Matrix LeapfrogScalingFactor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AuxiliaryTuningParameters -> Values
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VB.convert

hTuningParametersToSettings ::
  HSettings a ->
  TuningParameter ->
  AuxiliaryTuningParameters ->
  Either String (HSettings a)
hTuningParametersToSettings :: HSettings a
-> LeapfrogScalingFactor
-> AuxiliaryTuningParameters
-> Either String (HSettings a)
hTuningParametersToSettings HSettings a
s LeapfrogScalingFactor
t AuxiliaryTuningParameters
ts
  | Bool
nTsNotOK =
      String -> Either String (HSettings a)
forall a b. a -> Either a b
Left String
"hTuningParametersToSettings: Auxiliary variables do not have correct dimension."
  | Bool
otherwise =
      HSettings a -> Either String (HSettings a)
forall a b. b -> Either a b
Right (HSettings a -> Either String (HSettings a))
-> HSettings a -> Either String (HSettings a)
forall a b. (a -> b) -> a -> b
$
        HSettings a
s
          { hMasses :: Masses
hMasses = Masses
msTuned,
            hLeapfrogTrajectoryLength :: Int
hLeapfrogTrajectoryLength = Int
lTuned,
            hLeapfrogScalingFactor :: LeapfrogScalingFactor
hLeapfrogScalingFactor = LeapfrogScalingFactor
eTuned
          }
  where
    ms :: Masses
ms = HSettings a -> Masses
forall a. HSettings a -> Masses
hMasses HSettings a
s
    d :: Int
d = Matrix LeapfrogScalingFactor -> Int
forall t. Matrix t -> Int
L.rows (Matrix LeapfrogScalingFactor -> Int)
-> Matrix LeapfrogScalingFactor -> Int
forall a b. (a -> b) -> a -> b
$ Masses -> Matrix LeapfrogScalingFactor
forall t. Herm t -> Matrix t
L.unSym Masses
ms
    l :: Int
l = HSettings a -> Int
forall a. HSettings a -> Int
hLeapfrogTrajectoryLength HSettings a
s
    e :: LeapfrogScalingFactor
e = HSettings a -> LeapfrogScalingFactor
forall a. HSettings a -> LeapfrogScalingFactor
hLeapfrogScalingFactor HSettings a
s
    (HTune HTuneLeapfrog
tlf HTuneMasses
tms) = HSettings a -> HTune
forall a. HSettings a -> HTune
hTune HSettings a
s
    nTsNotOK :: Bool
nTsNotOK =
      let nTs :: Int
nTs = AuxiliaryTuningParameters -> Int
forall a. Unbox a => Vector a -> Int
VU.length AuxiliaryTuningParameters
ts
       in case HTuneMasses
tms of
            HTuneMasses
HNoTuneMasses -> Int
nTs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0
            HTuneMasses
_ -> Int
nTs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
d
    msTuned :: Masses
msTuned = case HTuneMasses
tms of
      HTuneMasses
HNoTuneMasses -> Masses
ms
      HTuneMasses
_ -> Int -> AuxiliaryTuningParameters -> Masses
tuningParametersToMasses Int
d AuxiliaryTuningParameters
ts
    -- The larger epsilon, the larger the proposal step size and the lower the
    -- expected acceptance ratio.
    --
    -- Further, we roughly keep \( L * \epsilon = 1.0 \). The equation is not
    -- correct, because we pull L closer to the original value to keep the
    -- runtime somewhat acceptable.
    (Int
lTuned, LeapfrogScalingFactor
eTuned) = case HTuneLeapfrog
tlf of
      HTuneLeapfrog
HNoTuneLeapfrog -> (Int
l, LeapfrogScalingFactor
e)
      HTuneLeapfrog
HTuneLeapfrog -> (LeapfrogScalingFactor -> Int
forall a b. (RealFrac a, Integral b) => a -> b
ceiling (LeapfrogScalingFactor -> Int) -> LeapfrogScalingFactor -> Int
forall a b. (a -> b) -> a -> b
$ Int -> LeapfrogScalingFactor
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
l LeapfrogScalingFactor
-> LeapfrogScalingFactor -> LeapfrogScalingFactor
forall a. Fractional a => a -> a -> a
/ (LeapfrogScalingFactor
t LeapfrogScalingFactor
-> LeapfrogScalingFactor -> LeapfrogScalingFactor
forall a. Floating a => a -> a -> a
** LeapfrogScalingFactor
0.9) :: Int, LeapfrogScalingFactor
t LeapfrogScalingFactor
-> LeapfrogScalingFactor -> LeapfrogScalingFactor
forall a. Num a => a -> a -> a
* LeapfrogScalingFactor
e)

hamiltonianSimpleWithTuningParameters ::
  HSettings a ->
  TuningParameter ->
  AuxiliaryTuningParameters ->
  Either String (ProposalSimple a)
hamiltonianSimpleWithTuningParameters :: HSettings a
-> LeapfrogScalingFactor
-> AuxiliaryTuningParameters
-> Either String (ProposalSimple a)
hamiltonianSimpleWithTuningParameters HSettings a
s LeapfrogScalingFactor
t AuxiliaryTuningParameters
ts =
  HSettings a
-> a
-> Gen RealWorld
-> IO (a, Log LeapfrogScalingFactor, Log LeapfrogScalingFactor)
forall a. HSettings a -> ProposalSimple a
hamiltonianSimple (HSettings a
 -> a
 -> Gen RealWorld
 -> IO (a, Log LeapfrogScalingFactor, Log LeapfrogScalingFactor))
-> Either String (HSettings a)
-> Either
     String
     (a
      -> Gen RealWorld
      -> IO (a, Log LeapfrogScalingFactor, Log LeapfrogScalingFactor))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HSettings a
-> LeapfrogScalingFactor
-> AuxiliaryTuningParameters
-> Either String (HSettings a)
forall a.
HSettings a
-> LeapfrogScalingFactor
-> AuxiliaryTuningParameters
-> Either String (HSettings a)
hTuningParametersToSettings HSettings a
s LeapfrogScalingFactor
t AuxiliaryTuningParameters
ts

-- The inverted covariance matrix and the log determinant of the covariance
-- matrix are calculated by 'hamiltonianSimple'.
hamiltonianSimpleWithMemoizedCovariance ::
  HSettings a ->
  HData ->
  ProposalSimple a
hamiltonianSimpleWithMemoizedCovariance :: HSettings a -> HData -> ProposalSimple a
hamiltonianSimpleWithMemoizedCovariance HSettings a
st HData
dt a
x GenIO
g = do
  Values
phi <- Values -> Masses -> GenIO -> IO Values
generateMomenta Values
mu Masses
masses GenIO
g
  Int
lRan <- (Int, Int) -> GenIO -> IO Int
forall a (m :: * -> *).
(Variate a, PrimMonad m) =>
(a, a) -> Gen (PrimState m) -> m a
uniformR (Int
lL, Int
lR) GenIO
g
  LeapfrogScalingFactor
eRan <- (LeapfrogScalingFactor, LeapfrogScalingFactor)
-> GenIO -> IO LeapfrogScalingFactor
forall a (m :: * -> *).
(Variate a, PrimMonad m) =>
(a, a) -> Gen (PrimState m) -> m a
uniformR (LeapfrogScalingFactor
eL, LeapfrogScalingFactor
eR) GenIO
g
  case (Values -> Values)
-> Maybe (Validate Values)
-> Masses
-> Int
-> LeapfrogScalingFactor
-> Values
-> Values
-> Maybe (Values, Values)
leapfrog Values -> Values
gradientVec Maybe (Validate Values)
mValVec Masses
massesInvEps Int
lRan LeapfrogScalingFactor
eRan Values
theta Values
phi of
    Maybe (Values, Values)
Nothing -> (a, Log LeapfrogScalingFactor, Log LeapfrogScalingFactor)
-> IO (a, Log LeapfrogScalingFactor, Log LeapfrogScalingFactor)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
x, Log LeapfrogScalingFactor
0.0, Log LeapfrogScalingFactor
1.0)
    Just (Values
theta', Values
phi') ->
      let -- Prior of momenta.
          prPhi :: Log LeapfrogScalingFactor
prPhi = Values
-> Masses
-> LeapfrogScalingFactor
-> Values
-> Log LeapfrogScalingFactor
logDensityMultivariateNormal Values
mu Masses
massesInv LeapfrogScalingFactor
logDetMasses Values
phi
          -- NOTE: Neal page 12: In order for the proposal to be in detailed
          -- balance, the momenta have to be negated before proposing the new
          -- value. This is not required here since the prior involves a
          -- multivariate normal distribution with means 0.
          prPhi' :: Log LeapfrogScalingFactor
prPhi' = Values
-> Masses
-> LeapfrogScalingFactor
-> Values
-> Log LeapfrogScalingFactor
logDensityMultivariateNormal Values
mu Masses
massesInv LeapfrogScalingFactor
logDetMasses Values
phi'
          kernelR :: Log LeapfrogScalingFactor
kernelR = Log LeapfrogScalingFactor
prPhi' Log LeapfrogScalingFactor
-> Log LeapfrogScalingFactor -> Log LeapfrogScalingFactor
forall a. Fractional a => a -> a -> a
/ Log LeapfrogScalingFactor
prPhi
       in (a, Log LeapfrogScalingFactor, Log LeapfrogScalingFactor)
-> IO (a, Log LeapfrogScalingFactor, Log LeapfrogScalingFactor)
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> Values -> a
fromVec a
x Values
theta', Log LeapfrogScalingFactor
kernelR, Log LeapfrogScalingFactor
1.0)
  where
    (HSettings a -> Values
toVec a -> Values -> a
fromVec Gradient a
gradient Maybe (Validate a)
mVal Masses
masses Int
l LeapfrogScalingFactor
e HTune
_) = HSettings a
st
    theta :: Values
theta = a -> Values
toVec a
x
    lL :: Int
lL = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [Int
1 :: Int, LeapfrogScalingFactor -> Int
forall a b. (RealFrac a, Integral b) => a -> b
floor (LeapfrogScalingFactor -> Int) -> LeapfrogScalingFactor -> Int
forall a b. (a -> b) -> a -> b
$ (LeapfrogScalingFactor
0.8 :: Double) LeapfrogScalingFactor
-> LeapfrogScalingFactor -> LeapfrogScalingFactor
forall a. Num a => a -> a -> a
* Int -> LeapfrogScalingFactor
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
l]
    lR :: Int
lR = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [Int
lL, LeapfrogScalingFactor -> Int
forall a b. (RealFrac a, Integral b) => a -> b
ceiling (LeapfrogScalingFactor -> Int) -> LeapfrogScalingFactor -> Int
forall a b. (a -> b) -> a -> b
$ (LeapfrogScalingFactor
1.2 :: Double) LeapfrogScalingFactor
-> LeapfrogScalingFactor -> LeapfrogScalingFactor
forall a. Num a => a -> a -> a
* Int -> LeapfrogScalingFactor
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
l]
    eL :: LeapfrogScalingFactor
eL = LeapfrogScalingFactor
0.8 LeapfrogScalingFactor
-> LeapfrogScalingFactor -> LeapfrogScalingFactor
forall a. Num a => a -> a -> a
* LeapfrogScalingFactor
e
    eR :: LeapfrogScalingFactor
eR = LeapfrogScalingFactor
1.2 LeapfrogScalingFactor
-> LeapfrogScalingFactor -> LeapfrogScalingFactor
forall a. Num a => a -> a -> a
* LeapfrogScalingFactor
e
    (HData Values
mu Masses
massesInv Masses
massesInvEps LeapfrogScalingFactor
logDetMasses) = HData
dt
    -- Vectorize the gradient and validation functions.
    gradientVec :: Values -> Values
gradientVec = a -> Values
toVec (a -> Values) -> (Values -> a) -> Values -> Values
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Gradient a
gradient Gradient a -> (Values -> a) -> Values -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Values -> a
fromVec a
x
    mValVec :: Maybe (Validate Values)
mValVec = Maybe (Validate a)
mVal Maybe (Validate a)
-> (Validate a -> Maybe (Validate Values))
-> Maybe (Validate Values)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (\Validate a
f -> Validate Values -> Maybe (Validate Values)
forall (m :: * -> *) a. Monad m => a -> m a
return (Validate Values -> Maybe (Validate Values))
-> Validate Values -> Maybe (Validate Values)
forall a b. (a -> b) -> a -> b
$ Validate a
f Validate a -> (Values -> a) -> Validate Values
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Values -> a
fromVec a
x)

hamiltonianSimple ::
  HSettings a ->
  ProposalSimple a
hamiltonianSimple :: HSettings a -> ProposalSimple a
hamiltonianSimple HSettings a
s = HSettings a -> HData -> ProposalSimple a
forall a. HSettings a -> HData -> ProposalSimple a
hamiltonianSimpleWithMemoizedCovariance HSettings a
s HData
hd
  where
    hd :: HData
hd = HSettings a -> HData
forall a. HSettings a -> HData
getHData HSettings a
s

-- If changed, also change help text of 'HTuneMasses'.
massMin :: Double
massMin :: LeapfrogScalingFactor
massMin = LeapfrogScalingFactor
1e-6

-- If changed, also change help text of 'HTuneMasses'.
massMax :: Double
massMax :: LeapfrogScalingFactor
massMax = LeapfrogScalingFactor
1e6

-- Minimal number of unique samples required for tuning the diagonal entries of
-- the mass matrix.
--
-- If changed, also change help text of 'HTuneMasses'.
samplesMinDiagonal :: Int
samplesMinDiagonal :: Int
samplesMinDiagonal = Int
61

-- Minimal number of samples required for tuning all entries of the mass matrix.
--
-- If changed, also change help text of 'HTuneMasses'.
samplesMinAll :: Int
samplesMinAll :: Int
samplesMinAll = Int
201

getSampleSize :: VS.Vector Double -> Int
getSampleSize :: Values -> Int
getSampleSize = Values -> Int
forall a. Storable a => Vector a -> Int
VS.length (Values -> Int) -> (Values -> Values) -> Values -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Values -> Values
forall a. (Storable a, Eq a) => Vector a -> Vector a
VS.uniq (Values -> Values) -> (Values -> Values) -> Values -> Values
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Values -> Values
forall e (v :: * -> *). (Ord e, Vector v e) => v e -> v e
S.gsort

-- Diagonal elements are variances which are strictly positive.
getNewMassDiagonalWithRescue :: Int -> Double -> Double -> Double
getNewMassDiagonalWithRescue :: Int
-> LeapfrogScalingFactor
-> LeapfrogScalingFactor
-> LeapfrogScalingFactor
getNewMassDiagonalWithRescue Int
sampleSize LeapfrogScalingFactor
massOld LeapfrogScalingFactor
massEstimate
  | Int
sampleSize Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
samplesMinDiagonal = LeapfrogScalingFactor
massOld
  -- NaN and negative masses could be errors.
  | LeapfrogScalingFactor -> Bool
forall a. RealFloat a => a -> Bool
isNaN LeapfrogScalingFactor
massEstimate = LeapfrogScalingFactor
massOld
  | LeapfrogScalingFactor
massEstimate LeapfrogScalingFactor -> LeapfrogScalingFactor -> Bool
forall a. Ord a => a -> a -> Bool
<= LeapfrogScalingFactor
0 = LeapfrogScalingFactor
massOld
  | LeapfrogScalingFactor
massMin LeapfrogScalingFactor -> LeapfrogScalingFactor -> Bool
forall a. Ord a => a -> a -> Bool
> LeapfrogScalingFactor
massNew = LeapfrogScalingFactor
massMin
  | LeapfrogScalingFactor
massNew LeapfrogScalingFactor -> LeapfrogScalingFactor -> Bool
forall a. Ord a => a -> a -> Bool
> LeapfrogScalingFactor
massMax = LeapfrogScalingFactor
massMax
  | Bool
otherwise = LeapfrogScalingFactor
massNew
  where
    massNewSqrt :: LeapfrogScalingFactor
massNewSqrt = LeapfrogScalingFactor -> LeapfrogScalingFactor
forall a. Fractional a => a -> a
recip LeapfrogScalingFactor
3 LeapfrogScalingFactor
-> LeapfrogScalingFactor -> LeapfrogScalingFactor
forall a. Num a => a -> a -> a
* (LeapfrogScalingFactor -> LeapfrogScalingFactor
forall a. Floating a => a -> a
sqrt LeapfrogScalingFactor
massOld LeapfrogScalingFactor
-> LeapfrogScalingFactor -> LeapfrogScalingFactor
forall a. Num a => a -> a -> a
+ LeapfrogScalingFactor
2 LeapfrogScalingFactor
-> LeapfrogScalingFactor -> LeapfrogScalingFactor
forall a. Num a => a -> a -> a
* LeapfrogScalingFactor -> LeapfrogScalingFactor
forall a. Floating a => a -> a
sqrt LeapfrogScalingFactor
massEstimate)
    massNew :: LeapfrogScalingFactor
massNew = LeapfrogScalingFactor
massNewSqrt LeapfrogScalingFactor
-> LeapfrogScalingFactor -> LeapfrogScalingFactor
forall a. Floating a => a -> a -> a
** LeapfrogScalingFactor
2

-- NOTE: Here, we lose time because we convert the states to vectors again,
-- something that has already been done. But then, auto tuning is not a runtime
-- determining factor.
tuneDiagonalMassesOnly ::
  Int ->
  (a -> Positions) ->
  AuxiliaryTuningFunction a
tuneDiagonalMassesOnly :: Int -> (a -> Values) -> AuxiliaryTuningFunction a
tuneDiagonalMassesOnly Int
dim a -> Values
toVec Vector a
xs AuxiliaryTuningParameters
ts
  -- If not enough data is available, do not tune.
  | Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
xs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
samplesMinDiagonal = AuxiliaryTuningParameters
ts
  | Bool
otherwise =
      -- Replace the diagonal.
      Masses -> AuxiliaryTuningParameters
massesToTuningParameters (Masses -> AuxiliaryTuningParameters)
-> Masses -> AuxiliaryTuningParameters
forall a b. (a -> b) -> a -> b
$
        Matrix LeapfrogScalingFactor -> Masses
forall t. Matrix t -> Herm t
L.trustSym (Matrix LeapfrogScalingFactor -> Masses)
-> Matrix LeapfrogScalingFactor -> Masses
forall a b. (a -> b) -> a -> b
$ Matrix LeapfrogScalingFactor
massesOld Matrix LeapfrogScalingFactor
-> Matrix LeapfrogScalingFactor -> Matrix LeapfrogScalingFactor
forall a. Num a => a -> a -> a
- Values -> Matrix LeapfrogScalingFactor
forall a. (Num a, Element a) => Vector a -> Matrix a
L.diag Values
massesDiagonalOld Matrix LeapfrogScalingFactor
-> Matrix LeapfrogScalingFactor -> Matrix LeapfrogScalingFactor
forall a. Num a => a -> a -> a
+ Values -> Matrix LeapfrogScalingFactor
forall a. (Num a, Element a) => Vector a -> Matrix a
L.diag Values
massesDiagonalNew
  where
    -- xs: Each vector entry contains all parameter values of one iteration.
    -- xs': Each row contains all parameter values of one iteration.
    xs' :: Matrix LeapfrogScalingFactor
xs' = [Values] -> Matrix LeapfrogScalingFactor
forall t. Element t => [Vector t] -> Matrix t
L.fromRows ([Values] -> Matrix LeapfrogScalingFactor)
-> [Values] -> Matrix LeapfrogScalingFactor
forall a b. (a -> b) -> a -> b
$ Vector Values -> [Values]
forall a. Vector a -> [a]
VB.toList (Vector Values -> [Values]) -> Vector Values -> [Values]
forall a b. (a -> b) -> a -> b
$ (a -> Values) -> Vector a -> Vector Values
forall a b. (a -> b) -> Vector a -> Vector b
VB.map a -> Values
toVec Vector a
xs
    sampleSizes :: Vector Int
sampleSizes = [Int] -> Vector Int
forall a. Storable a => [a] -> Vector a
VS.fromList ([Int] -> Vector Int) -> [Int] -> Vector Int
forall a b. (a -> b) -> a -> b
$ (Values -> Int) -> [Values] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Values -> Int
getSampleSize ([Values] -> [Int]) -> [Values] -> [Int]
forall a b. (a -> b) -> a -> b
$ Matrix LeapfrogScalingFactor -> [Values]
forall t. Element t => Matrix t -> [Vector t]
L.toColumns Matrix LeapfrogScalingFactor
xs'
    massesOld :: Matrix LeapfrogScalingFactor
massesOld = Masses -> Matrix LeapfrogScalingFactor
forall t. Herm t -> Matrix t
L.unSym (Masses -> Matrix LeapfrogScalingFactor)
-> Masses -> Matrix LeapfrogScalingFactor
forall a b. (a -> b) -> a -> b
$ Int -> AuxiliaryTuningParameters -> Masses
tuningParametersToMasses Int
dim AuxiliaryTuningParameters
ts
    massesDiagonalOld :: Values
massesDiagonalOld = Matrix LeapfrogScalingFactor -> Values
forall t. Element t => Matrix t -> Vector t
L.takeDiag Matrix LeapfrogScalingFactor
massesOld
    massesDiagonalEstimate :: Values
massesDiagonalEstimate = [LeapfrogScalingFactor] -> Values
forall a. Storable a => [a] -> Vector a
VS.fromList ([LeapfrogScalingFactor] -> Values)
-> [LeapfrogScalingFactor] -> Values
forall a b. (a -> b) -> a -> b
$ (Values -> LeapfrogScalingFactor)
-> [Values] -> [LeapfrogScalingFactor]
forall a b. (a -> b) -> [a] -> [b]
map (LeapfrogScalingFactor -> LeapfrogScalingFactor
forall a. Fractional a => a -> a
recip (LeapfrogScalingFactor -> LeapfrogScalingFactor)
-> (Values -> LeapfrogScalingFactor)
-> Values
-> LeapfrogScalingFactor
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Values -> LeapfrogScalingFactor
forall (v :: * -> *).
Vector v LeapfrogScalingFactor =>
v LeapfrogScalingFactor -> LeapfrogScalingFactor
S.variance) ([Values] -> [LeapfrogScalingFactor])
-> [Values] -> [LeapfrogScalingFactor]
forall a b. (a -> b) -> a -> b
$ Matrix LeapfrogScalingFactor -> [Values]
forall t. Element t => Matrix t -> [Vector t]
L.toColumns Matrix LeapfrogScalingFactor
xs'
    massesDiagonalNew :: Values
massesDiagonalNew =
      (Int
 -> LeapfrogScalingFactor
 -> LeapfrogScalingFactor
 -> LeapfrogScalingFactor)
-> Vector Int -> Values -> Values -> Values
forall a b c d.
(Storable a, Storable b, Storable c, Storable d) =>
(a -> b -> c -> d) -> Vector a -> Vector b -> Vector c -> Vector d
VS.zipWith3
        Int
-> LeapfrogScalingFactor
-> LeapfrogScalingFactor
-> LeapfrogScalingFactor
getNewMassDiagonalWithRescue
        Vector Int
sampleSizes
        Values
massesDiagonalOld
        Values
massesDiagonalEstimate

-- NOTE: Here, we lose time because we convert the states to vectors again,
-- something that has already been done. But then, auto tuning is not a runtime
-- determining factor.
tuneAllMasses ::
  Int ->
  (a -> Positions) ->
  AuxiliaryTuningFunction a
tuneAllMasses :: Int -> (a -> Values) -> AuxiliaryTuningFunction a
tuneAllMasses Int
dim a -> Values
toVec Vector a
xs AuxiliaryTuningParameters
ts
  -- If not enough data is available, do not tune.
  | Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
xs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
samplesMinDiagonal = AuxiliaryTuningParameters
ts
  -- If not enough data is available, only the diagonal masses are tuned.
  | Vector a -> Int
forall a. Vector a -> Int
VB.length Vector a
xs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
samplesMinAll = AuxiliaryTuningParameters
fallbackDiagonal
  | Matrix LeapfrogScalingFactor -> Int
forall t. Field t => Matrix t -> Int
L.rank Matrix LeapfrogScalingFactor
xs' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
dim = AuxiliaryTuningParameters
fallbackDiagonal
  | Bool
otherwise = Masses -> AuxiliaryTuningParameters
massesToTuningParameters (Masses -> AuxiliaryTuningParameters)
-> Masses -> AuxiliaryTuningParameters
forall a b. (a -> b) -> a -> b
$ Matrix LeapfrogScalingFactor -> Masses
forall t. Matrix t -> Herm t
L.trustSym Matrix LeapfrogScalingFactor
massesNew
  where
    fallbackDiagonal :: AuxiliaryTuningParameters
fallbackDiagonal = Int -> (a -> Values) -> AuxiliaryTuningFunction a
forall a. Int -> (a -> Values) -> AuxiliaryTuningFunction a
tuneDiagonalMassesOnly Int
dim a -> Values
toVec Vector a
xs AuxiliaryTuningParameters
ts
    -- xs: Each vector entry contains all parameter values of one iteration.
    -- xs': Each row contains all parameter values of one iteration.
    xs' :: Matrix LeapfrogScalingFactor
xs' = [Values] -> Matrix LeapfrogScalingFactor
forall t. Element t => [Vector t] -> Matrix t
L.fromRows ([Values] -> Matrix LeapfrogScalingFactor)
-> [Values] -> Matrix LeapfrogScalingFactor
forall a b. (a -> b) -> a -> b
$ Vector Values -> [Values]
forall a. Vector a -> [a]
VB.toList (Vector Values -> [Values]) -> Vector Values -> [Values]
forall a b. (a -> b) -> a -> b
$ (a -> Values) -> Vector a -> Vector Values
forall a b. (a -> b) -> Vector a -> Vector b
VB.map a -> Values
toVec Vector a
xs
    (Values
_, Values
ss, Matrix LeapfrogScalingFactor
xsNormalized) = Matrix LeapfrogScalingFactor
-> (Values, Values, Matrix LeapfrogScalingFactor)
S.scale Matrix LeapfrogScalingFactor
xs'
    -- sigmaNormalized = L.unSym $ either error id $ S.oracleApproximatingShrinkage xsNormalized
    sigmaNormalized :: Matrix LeapfrogScalingFactor
sigmaNormalized = Masses -> Matrix LeapfrogScalingFactor
forall t. Herm t -> Matrix t
L.unSym (Masses -> Matrix LeapfrogScalingFactor)
-> Masses -> Matrix LeapfrogScalingFactor
forall a b. (a -> b) -> a -> b
$ (String -> Masses)
-> ((Masses, Masses) -> Masses)
-> Either String (Masses, Masses)
-> Masses
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either String -> Masses
forall a. HasCallStack => String -> a
error (Masses, Masses) -> Masses
forall a b. (a, b) -> a
fst (Either String (Masses, Masses) -> Masses)
-> Either String (Masses, Masses) -> Masses
forall a b. (a -> b) -> a -> b
$ LeapfrogScalingFactor
-> Matrix LeapfrogScalingFactor -> Either String (Masses, Masses)
S.graphicalLasso LeapfrogScalingFactor
0.5 Matrix LeapfrogScalingFactor
xsNormalized
    sigma :: Matrix LeapfrogScalingFactor
sigma = Values
-> Matrix LeapfrogScalingFactor -> Matrix LeapfrogScalingFactor
S.rescaleSWith Values
ss Matrix LeapfrogScalingFactor
sigmaNormalized
    massesNew :: Matrix LeapfrogScalingFactor
massesNew = Matrix LeapfrogScalingFactor -> Matrix LeapfrogScalingFactor
forall t. Field t => Matrix t -> Matrix t
L.inv Matrix LeapfrogScalingFactor
sigma

-- | Hamiltonian Monte Carlo proposal.
hamiltonian ::
  Eq a =>
  -- | The sample state is used for error checks and to calculate the dimension
  -- of the proposal.
  a ->
  HSettings a ->
  PName ->
  PWeight ->
  Proposal a
hamiltonian :: a -> HSettings a -> PName -> PWeight -> Proposal a
hamiltonian a
x HSettings a
s PName
n PWeight
w = case a -> HSettings a -> Maybe String
forall a. Eq a => a -> HSettings a -> Maybe String
checkHSettings a
x HSettings a
s of
  Just String
err -> String -> Proposal a
forall a. HasCallStack => String -> a
error String
err
  Maybe String
Nothing ->
    let desc :: PDescription
desc = String -> PDescription
PDescription String
"Hamiltonian Monte Carlo (HMC)"
        toVec :: a -> Values
toVec = HSettings a -> a -> Values
forall a. HSettings a -> a -> Values
hToVector HSettings a
s
        dim :: IndexOf Vector
dim = (Values -> IndexOf Vector
forall (c :: * -> *) t. Container c t => c t -> IndexOf c
L.size (Values -> IndexOf Vector) -> Values -> IndexOf Vector
forall a b. (a -> b) -> a -> b
$ a -> Values
toVec a
x)
        pDim :: PDimension
pDim = Int -> LeapfrogScalingFactor -> PDimension
PSpecial Int
dim LeapfrogScalingFactor
0.65
        ts :: AuxiliaryTuningParameters
ts = Masses -> AuxiliaryTuningParameters
massesToTuningParameters (HSettings a -> Masses
forall a. HSettings a -> Masses
hMasses HSettings a
s)
        ps :: ProposalSimple a
ps = HSettings a -> ProposalSimple a
forall a. HSettings a -> ProposalSimple a
hamiltonianSimple HSettings a
s
        hamiltonianWith :: Maybe (Tuner a) -> Proposal a
hamiltonianWith = PName
-> PDescription
-> PDimension
-> PWeight
-> ProposalSimple a
-> Maybe (Tuner a)
-> Proposal a
forall a.
PName
-> PDescription
-> PDimension
-> PWeight
-> ProposalSimple a
-> Maybe (Tuner a)
-> Proposal a
Proposal PName
n PDescription
desc PDimension
pDim PWeight
w a
-> Gen RealWorld
-> IO (a, Log LeapfrogScalingFactor, Log LeapfrogScalingFactor)
ProposalSimple a
ps
        tSet :: HTune
tSet@(HTune HTuneLeapfrog
tlf HTuneMasses
tms) = HSettings a -> HTune
forall a. HSettings a -> HTune
hTune HSettings a
s
        tFun :: LeapfrogScalingFactor
-> LeapfrogScalingFactor -> LeapfrogScalingFactor
tFun = case HTuneLeapfrog
tlf of
          HTuneLeapfrog
HNoTuneLeapfrog -> LeapfrogScalingFactor
-> LeapfrogScalingFactor -> LeapfrogScalingFactor
noTuningFunction
          HTuneLeapfrog
HTuneLeapfrog -> PDimension
-> LeapfrogScalingFactor
-> LeapfrogScalingFactor
-> LeapfrogScalingFactor
defaultTuningFunctionWith PDimension
pDim
        tFunAux :: AuxiliaryTuningFunction a
tFunAux = case HTuneMasses
tms of
          HTuneMasses
HNoTuneMasses -> AuxiliaryTuningFunction a
forall a. AuxiliaryTuningFunction a
noAuxiliaryTuningFunction
          HTuneMasses
HTuneDiagonalMassesOnly -> Int -> (a -> Values) -> AuxiliaryTuningFunction a
forall a. Int -> (a -> Values) -> AuxiliaryTuningFunction a
tuneDiagonalMassesOnly Int
dim a -> Values
toVec
          HTuneMasses
HTuneAllMasses -> Int -> (a -> Values) -> AuxiliaryTuningFunction a
forall a. Int -> (a -> Values) -> AuxiliaryTuningFunction a
tuneAllMasses Int
dim a -> Values
toVec
     in case HTune
tSet of
          (HTune HTuneLeapfrog
HNoTuneLeapfrog HTuneMasses
HNoTuneMasses) -> Maybe (Tuner a) -> Proposal a
hamiltonianWith Maybe (Tuner a)
forall a. Maybe a
Nothing
          HTune
_ ->
            let tuner :: Tuner a
tuner = LeapfrogScalingFactor
-> (LeapfrogScalingFactor
    -> LeapfrogScalingFactor -> LeapfrogScalingFactor)
-> AuxiliaryTuningParameters
-> AuxiliaryTuningFunction a
-> (LeapfrogScalingFactor
    -> AuxiliaryTuningParameters -> Either String (ProposalSimple a))
-> Tuner a
forall a.
LeapfrogScalingFactor
-> (LeapfrogScalingFactor
    -> LeapfrogScalingFactor -> LeapfrogScalingFactor)
-> AuxiliaryTuningParameters
-> AuxiliaryTuningFunction a
-> (LeapfrogScalingFactor
    -> AuxiliaryTuningParameters -> Either String (ProposalSimple a))
-> Tuner a
Tuner LeapfrogScalingFactor
1.0 LeapfrogScalingFactor
-> LeapfrogScalingFactor -> LeapfrogScalingFactor
tFun AuxiliaryTuningParameters
ts AuxiliaryTuningFunction a
tFunAux (HSettings a
-> LeapfrogScalingFactor
-> AuxiliaryTuningParameters
-> Either String (ProposalSimple a)
forall a.
HSettings a
-> LeapfrogScalingFactor
-> AuxiliaryTuningParameters
-> Either String (ProposalSimple a)
hamiltonianSimpleWithTuningParameters HSettings a
s)
             in Maybe (Tuner a) -> Proposal a
hamiltonianWith (Maybe (Tuner a) -> Proposal a) -> Maybe (Tuner a) -> Proposal a
forall a b. (a -> b) -> a -> b
$ Tuner a -> Maybe (Tuner a)
forall a. a -> Maybe a
Just Tuner a
tuner