diff --git a/Numeric/MCMC/AffineInvariantEnsemble.hs b/Numeric/MCMC/AffineInvariantEnsemble.hs --- a/Numeric/MCMC/AffineInvariantEnsemble.hs +++ b/Numeric/MCMC/AffineInvariantEnsemble.hs @@ -1,20 +1,24 @@ -{-# LANGUAGE FlexibleInstances, FlexibleContexts #-} +{-# 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 highly skewed or anisotropic distributions. +-- 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 for details --- of the general sampling routine. +-- of the general sampling routine. module Numeric.MCMC.AffineInvariantEnsemble ( -- * Data structures - Config(..), AffineTransform(..), Trace + Config(..), AffineTransform(..), Trace(..) -- * Chain management - , runChain, initializeEnsemble, defaultSeed, prune + , runChain, initializeEnsemble, defaultSeed + -- * Trace management + , prune, ptrace ) where import Numeric.MCMC.Util -import Data.List (foldl') -import Data.List.Split (splitEvery) +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 @@ -31,23 +35,31 @@ libError :: String libError = "Numeric.MCMC.AffineInvariantEnsemble." --- | A data type holding the configuration of the Markov chain at any given epoch. `ensemble` accesses --- the IntMap constituting the current ensemble, while 'accepts' records the number of proposals that --- have been accepted up to the current epoch. +-- | 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 data type representing the affine transformation to be used on particles in 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) - --- | A data type holding a chain's trace. -newtype Trace a = Trace [[a]] +-- | A chain's trace. +newtype Trace = Trace [[Double]] -instance Show (Trace Double) where +-- | 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 @@ -55,31 +67,32 @@ -> Double -- ^ Random double drawn from appropriate distribution -> ([Double] -> Double) -- ^ Target function -> ([Double], Double) -- ^ Tuple containing proposed move and its log acceptance prob -stretch xk xj nw z target = (proposal, logAP) - where proposal = zipWith (+) (map (*z) xk) (map (*(1-z)) xj) - logAP = let val = target proposal - target xk + (fromIntegral nw - 1) * log z - in if val > 0 then 0 else val +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 :: (Fractional c, Num t, Ord t) - => [c] -- ^ Focused walker - -> [[c]] -- ^ Sub-ensemble of n alternate walkers - -> [c] -- ^ n random doubles drawn from a standard normal - -> ([c] -> t) -- ^ Target function - -> ([c], t) -- ^ Tuple containing proposed move and its log acceptance prob -walk xk xjs zs target = let val = target proposal - target xk in (proposal, if val > 0 then 0 else val) - where nxjs = length xjs - xjsmean = map (/ fromIntegral nxjs) $ listReducer (length xk) xjs - xjscentd = zipWith (zipWith (-)) xjs (replicate nxjs xjsmean) - listReducer n = foldl' (zipWith (+)) (replicate n 0.0) - proposal = zipWith (+) xk (listReducer nxjs $ zipWith (\z -> map (*z)) zs xjscentd) +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 #-} --- | Naively initialize an ensemble. Creates a 'Config' containing /nw/ walkers, each of dimension /nd/, +-- | 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 this is expected to be a region of low density, you'll probably want to specify +-- 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." @@ -88,7 +101,7 @@ | otherwise = do gen <- create inits <- replicateM (nw * nd) (uniformR (0 :: Double, 1) gen) - let arr = IntMap.fromList $ zip [1..] (splitEvery nd inits) + let arr = IntMap.fromList $ zip [1..] (chunksOf nd inits) initConfig = Config {ensemble = arr, accepts = 0} return initConfig @@ -119,7 +132,8 @@ g0 <- restore seed altWalkerIndex <- genDiffInt targetWalkerIndex (1, numWalkers) g0 z0 <- uniformR (0, 1) g0 - let z = 0.5 * (z0 + 1) * (z0 + 1) + 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 @@ -132,7 +146,7 @@ where createEnsemble = map (\k -> fromJust (IntMap.lookup k walkers)) return $ walk targetWalker altWalkerEnsemble zs target - -- Compare and possible accept proposal + -- Compare and possibly accept proposal when (zc <= exp logAcceptanceProb) $ writeSTRef stConfig Config {ensemble = IntMap.update (\_ -> Just proposal) targetWalkerIndex walkers, accepts = nacc + 1} @@ -140,6 +154,7 @@ endConfig <- readSTRef stConfig let endPosition = ensemble endConfig return endPosition +{-# INLINE moveEnsemble #-} -- | Typical use: -- @@ -153,13 +168,12 @@ -- 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 function. --- --- The /target/ must be a function with type @[Double] -> Double@. Functions using more complicated data structures --- internally can simply be curried to this type. +-- 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. -- --- Examples of use can be found at . -runChain :: Vector v Word32 => Int -> ([Double] -> Double) -> Config -> v Word32 -> AffineTransform Int -> (Config, Trace Double) +-- The /target/ must be a function with type @[Double] -> Double@. Examples of use can be found at +-- . +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 @@ -172,10 +186,6 @@ results <- readSTRef config return (results, trace) - --- | Prune some initial epochs (i.e. suspected burn-in) from a 'Trace'. -prune :: Int -> Trace Double -> Trace Double -prune n (Trace xs) = Trace (drop n xs) -- | The default seed provided by the library. This seed is different from the one used internally in 'initializeEnsemble'. defaultSeed :: U.Vector Word32 diff --git a/Numeric/MCMC/Util.hs b/Numeric/MCMC/Util.hs --- a/Numeric/MCMC/Util.hs +++ b/Numeric/MCMC/Util.hs @@ -1,3 +1,5 @@ +{-# LANGUAGE BangPatterns #-} + -- | Various internal utilities. module Numeric.MCMC.Util @@ -11,30 +13,33 @@ import Control.Monad import Control.Monad.ST import Data.STRef +import Data.List.Split (chunksOf) -- | Given Int, bounds, and generator, generate a different Int in the bound. genDiffInt :: PrimMonad m => Int -> (Int, Int) -> Gen (PrimState m) -> m Int genDiffInt a bounds gen = fix $ \loopB -> do b <- uniformR bounds gen if a == b then loopB else return b +{-# INLINE genDiffInt #-} --- | Tail-recursive, list-fused mean function. -mean :: [Double] -> Double +-- | Tail-recursive mean function. +mean :: [Double] -> Double mean = go 0.0 0 where go :: Double -> Int -> [Double] -> Double - go s l [] = s / fromIntegral l - go s l (x:xs) = go (s + x) (l + 1) xs - --- | Map a function over a pair. -mapPair :: (a -> b) -> (a, a) -> (b, b) -mapPair f (a, b) = (f a, f b) + go !s !l [] = s / fromIntegral l + go !s !l (x:xs) = go (s + x) (l + 1) xs +{-# INLINE mean #-} --- | Convert a list to a pair. -shortListToPair :: [a] -> (a, a) -shortListToPair [x0, x1] = (x0, x1) -shortListToPair _ = error "shortListToPair - list must have length 2." +-- | Tail-recursive mean function for lists. +listMean :: [[Double]] -> [Double] +listMean = go (repeat 0) 0 + where go :: [Double] -> Int -> [[Double]] -> [Double] + go !s !l [] = map (/ fromIntegral l) s + go !s !l (xs:xss) = go (zipWith (+) s xs) (l + 1) xss +{-# INLINE listMean #-} --- | Knuth-shuffle a list. Uses 'Seq' internally. +-- | Knuth-shuffle a list. Uses 'Seq' internally. +-- FIXME Would be nice if it was lazy, to work well with `sample` shuffle :: [a] -> Seed -> [a] shuffle xs seed = runST $ do xsref <- newSTRef $ Seq.fromList xs @@ -51,9 +56,45 @@ result <- readSTRef xsref return $ toList result +{-# INLINE shuffle #-} -- | Sample from a list without replacement. sample :: Int -> [a] -> Seed -> [a] sample k xs seed = take k $ shuffle xs seed +{-# INLINE sample #-} +-- | Autocovariance function. `m` is the maximum lag. +acov :: [Double] -> Int -> [Double] +acov xs m = go [] m xc + where xc = take (l - m) $ map (subtract (mean xs)) xs + l = length xs + go !s0 !j0 ys | j0 < 0 = map (/ fromIntegral (l - m)) (reverse s0) + | otherwise = go (sum (zipWith (*) xc ys) : s0) (j0 - 1) (drop 1 ys) +{-# INLINE acov #-} + +-- | (mean, standard error, integrated autocorrelation time) +acor :: [Double] -> (Double, Double, Double) +acor xs = let t0 = d0 / head ac0 + in if t0*fromIntegral winmult < fromIntegral maxlag + then (mean xs, sqrt (d0 / fromIntegral l0), t0) + else (mean xs, sqrt (df / fromIntegral l0), df / head acf) + where + -- Initial values and constants + (l0, ac0, d0) = (length xs, acov xs maxlag, head ac0 + 2 * sum (tail ac0)) + (taumax, winmult) = (2, 5) :: (Int, Int) + maxlag = taumax*winmult + + -- Final autocovariance + (sig1, acf) = go xs ac0 (sqrt (d0 / fromIntegral l0)) (d0 / head ac0) + df = 0.25 * sig1 * sig1 * fromIntegral l0 + + -- The recursive worker + go ys ac sig !tau + | tau*fromIntegral winmult < fromIntegral maxlag = (sig, ac) + | otherwise = + let (ys1, ac1, d, t1) = (joiner ys, acov ys1 maxlag, head ac1 + (2 :: Double) * sum (tail ac1), d / head ac1) + in if ys1 == [] then (sig, ac) else go ys1 ac1 (sqrt (d / fromIntegral (length ys))) t1 + where + joiner zs = map sum $ take (truncate $ fromIntegral (length zs) / (2 :: Double)) (chunksOf 2 zs) +{-# INLINE acor #-} diff --git a/affine-invariant-ensemble-mcmc.cabal b/affine-invariant-ensemble-mcmc.cabal --- a/affine-invariant-ensemble-mcmc.cabal +++ b/affine-invariant-ensemble-mcmc.cabal @@ -1,6 +1,6 @@ Name: affine-invariant-ensemble-mcmc Homepage: http://github.com/jtobin/affine-invariant-ensemble-mcmc -Version: 0.1.0.0 +Version: 0.2.0.0 Cabal-version: >=1.8 build-type: Simple License: BSD3 @@ -11,14 +11,14 @@ Synopsis: General-purpose sampling Description: - A general-purpose sampling routine for badly-scaled distributions. + A general-purpose sampler for anisotropic distributions. Source-repository head Type: git Location: http://github.com/jtobin/affine-invariant-ensemble-mcmc.git Library - Build-depends: base ==4.5.*, containers ==0.5.*, vector ==0.9.*, mwc-random ==0.12.*, primitive ==0.4.*, split ==0.1.* + Build-depends: base ==4.5.*, containers ==0.5.*, vector ==0.9.*, mwc-random ==0.12.*, primitive ==0.4.*, split ==0.2.* ghc-options: -Wall Exposed-modules: Numeric.MCMC.AffineInvariantEnsemble, Numeric.MCMC.Util