{-# LANGUAGE FlexibleContexts, FlexibleInstances #-}
-- | A Haskell implementation of Goodman & Weare (2010)'s /affine invariant ensemble MCMC/, a family of Markov
-- Chain Monte Carlo methods that can efficiently sample from anisotropic distributions. This algorithm
-- should provide state-of-the-art sampling on continuous, roughly-unimodal distributions, independent of
-- correlations between parameters.
--
-- See 'runChain' for an overview of use, and <http://msp.berkeley.edu/camcos/2010/5-1/p04.xhtml> for details
-- of the general sampling routine.
module Numeric.MCMC.AffineInvariantEnsemble (
-- * Data structures
Config(..), AffineTransform(..), Trace(..)
-- * Chain management
, runChain, initializeEnsemble, defaultSeed
-- * Trace management
, prune, ptrace
) where
import Numeric.MCMC.Util
import Data.List (foldl', transpose)
import Data.List.Split (chunksOf)
import Data.IntMap.Strict (IntMap)
import qualified Data.IntMap.Strict as IntMap
import qualified Data.Vector.Unboxed as U
import Data.Vector.Generic (Vector)
import System.Random.MWC
import System.Random.MWC.Distributions (standard)
import Control.Monad.ST (ST, runST)
import Data.STRef (STRef, newSTRef, readSTRef, writeSTRef)
import Control.Monad (forM_, replicateM, when)
import Data.Maybe (fromJust)
import Control.Monad.Primitive (PrimMonad)
import Data.Word (Word32)
libError :: String
libError = "Numeric.MCMC.AffineInvariantEnsemble."
-- | The state of the Markov chain. `ensemble` accesses the current ensemble, while
-- 'accepts' accesses the number of proposals that have been accepted up to the current epoch.
data Config = Config { ensemble :: !(IntMap [Double])
, accepts :: {-# UNPACK #-} !Int
}
-- | A chain's trace.
newtype Trace = Trace [[Double]]
-- | A simple Show instance for 'Trace'.
instance Show Trace where
show (Trace xs) = unlines $ map (unwords . map show) xs
-- | Prune some initial epochs (i.e. suspected burn-in) from a 'Trace'.
prune :: Int -> Trace -> Trace
prune n (Trace xs) = Trace (drop n xs)
-- | Retrieve parameter number `n` from a 'Trace'.
ptrace :: Int -> Trace -> [Double]
ptrace n (Trace xs) = transpose xs !! n
-- | The affine transformation to be used to propose moves within an ensemble. The general-purpose
-- /stretch/ and /walk/ transformations described in Goodman and Weare (2010) are supported.
data AffineTransform a = Stretch | Walk a deriving (Eq, Read)
-- The `stretch` affine transform.
stretch :: [Double] -- ^ Focused walker
-> [Double] -- ^ Alternate walker
-> Int -- ^ Number of walkers in ensemble
-> Double -- ^ Random double drawn from appropriate distribution
-> ([Double] -> Double) -- ^ Target function
-> ([Double], Double) -- ^ Tuple containing proposed move and its log acceptance prob
stretch targetWalker altWalker nw z target =
let val = target proposal - target targetWalker + (fromIntegral nw - 1) * log z
in (proposal, if val > 0 then 0 else val)
where proposal = zipWith (+) (map (*z) targetWalker) (map (*(1-z)) altWalker)
{-# INLINE stretch #-}
-- The `walk` affine transform.
walk :: [Double] -- ^ Focused walker
-> [[Double]] -- ^ Sub-ensemble of n alternate walkers
-> [Double] -- ^ n random doubles drawn from a standard normal
-> ([Double] -> Double) -- ^ Target function
-> ([Double], Double) -- ^ Tuple containing proposed move and its log acceptance prob
walk targetWalker subEnsemble zs target =
let val = target proposal - target targetWalker
in (proposal, if val > 0 then 0 else val)
where subEnsembleMean = listMean subEnsemble
subEnsembleCentd = map (flip (zipWith (-)) subEnsembleMean) subEnsemble
proposal = zipWith (+) targetWalker (foldl' (zipWith (+)) (cycle [0]) (map (\(a, b) -> map (*a) b) $ zip zs subEnsembleCentd))
{-# INLINE walk #-}
-- | A convenience function to naively initialize an ensemble. Creates a 'Config' containing /nw/ walkers, each of dimension /nd/,
-- and initializes 'accepts' at 0. Each dimensional element is drawn randomly from (0,1] (using a different
-- seed than 'defaultSeed'). If the [0, 1] hypercube is expected to be a region of low density, you'll probably want to specify
-- your own initial configuration.
--
-- This function will run in either the ST or IO monads.
initializeEnsemble :: PrimMonad m => Int -> Int -> m Config
initializeEnsemble nw nd
| nw < 2 = error $ libError ++ "initializeEnsemble: Number of walkers must be >= 2."
| nd < 1 = error $ libError ++ "initializeEnsemble: Number of dimensions must be >= 1."
| nw < nd = error $ libError ++ "initializeEnsemble: Number of walkers should be greater than number of dimensions."
| otherwise = do
gen <- create
inits <- replicateM (nw * nd) (uniformR (0 :: Double, 1) gen)
let arr = IntMap.fromList $ zip [1..] (chunksOf nd inits)
initConfig = Config {ensemble = arr, accepts = 0}
return initConfig
-- Move an ensemble forward one step.
moveEnsemble :: Int -- Number of walkers in the ensemble.
-> STRef s Config -- A STRef storing the ensemble configuration.
-> Gen s -- Random number generator
-> ([Double] -> Double) -- Desired target
-> AffineTransform Int -- Affine transformation to use
-> ST s (IntMap [Double])
moveEnsemble numWalkers _ _ _ (Walk n)
| n >= numWalkers = error "Numeric.MCMC.AffineInvariantEnsemble moveEnsemble: size of `Walk` sub-ensemble must be strictly less than size of full ensemble."
moveEnsemble numWalkers stConfig gen target xform = do
forM_ [1..numWalkers] $ \targetWalkerIndex -> do
-- Inits
config <- readSTRef stConfig
let walkers = ensemble config
nacc = accepts config
targetWalker = fromJust $ IntMap.lookup targetWalkerIndex walkers
-- Randomness
zc <- uniformR (0, 1) gen
seed <- save gen
-- Generate proposal and acceptance probability
let (proposal, logAcceptanceProb) = case xform of
Stretch -> runST $ do
g0 <- restore seed
altWalkerIndex <- genDiffInt targetWalkerIndex (1, numWalkers) g0
z0 <- uniformR (0, 1) g0
let -- z = 0.5 * (z0 + 1) * (z0 + 1)
z = let a = 2 :: Double in ((a - 1)*z0 + 1)^(2 :: Int) / a
altWalker = fromJust $ IntMap.lookup altWalkerIndex walkers
return $ stretch targetWalker altWalker numWalkers z target
Walk n -> runST $ do
g0 <- restore seed
zs <- replicateM n (standard g0)
g1 <- save g0
let subMapKeys = sample n (IntMap.keys walkers) g1
altWalkerEnsemble = createEnsemble subMapKeys
where createEnsemble = map (\k -> fromJust (IntMap.lookup k walkers))
return $ walk targetWalker altWalkerEnsemble zs target
-- Compare and possibly accept proposal
when (zc <= exp logAcceptanceProb) $
writeSTRef stConfig Config {ensemble = IntMap.update (\_ -> Just proposal) targetWalkerIndex walkers, accepts = nacc + 1}
-- Return end state
endConfig <- readSTRef stConfig
let endPosition = ensemble endConfig
return endPosition
{-# INLINE moveEnsemble #-}
-- | Typical use:
--
-- @
-- runChain steps target initConfig seed xform
-- @
--
-- Run the Markov chain for /steps/ epochs. The chain will wander over /target/'s parameter space such that,
-- after \"long enough\", the points it visits will effectively be independent samples from the distribution
-- proportional to /target/. The Markov chain procedes by possibly applying an affine transformation to each of
-- the particles contained in 'ensemble' /initConfig/, sequentially.
--
-- This function will return a tuple contanining 1) the 'Config' corresponding to the final epoch of the chain,
-- and 2) the chain's 'Trace'. The 'Trace' can be used, for example, to approximate integrals of the target.
-- A `Show` instance exists for pretty-printing to stdout.
--
-- The /target/ must be a function with type @[Double] -> Double@. Examples of use can be found at
-- <http://github.com/jtobin/affine-invariant-ensemble-mcmc/tree/master/Numeric/MCMC/Examples>.
runChain :: Vector v Word32 => Int -> ([Double] -> Double) -> Config -> v Word32 -> AffineTransform Int -> (Config, Trace)
runChain steps target initConfig seed xform
| steps < 1 = error $ libError ++ "runChain: `steps` must be >= 1."
| otherwise = runST $ do
let nw = IntMap.size (ensemble initConfig)
gen <- initialize seed
config <- newSTRef initConfig
frames <- replicateM steps (moveEnsemble nw config gen target xform)
let trace = Trace (map snd $ concatMap IntMap.toList frames)
results <- readSTRef config
return (results, trace)
-- | The default seed provided by the library. This seed is different from the one used internally in 'initializeEnsemble'.
defaultSeed :: U.Vector Word32
defaultSeed = U.singleton 42