affine-invariant-ensemble-mcmc 0.1.0.0 → 0.2.0.0
raw patch · 3 files changed
+117/−66 lines, 3 filesdep ~split
Dependency ranges changed: split
Files
- Numeric/MCMC/AffineInvariantEnsemble.hs +60/−50
- Numeric/MCMC/Util.hs +54/−13
- affine-invariant-ensemble-mcmc.cabal +3/−3
Numeric/MCMC/AffineInvariantEnsemble.hs view
@@ -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 <http://msp.berkeley.edu/camcos/2010/5-1/p04.xhtml> 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 <http://github.com/jtobin/affine-invariant-ensemble-mcmc/Numeric/MCMC/Examples>.-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 +-- <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@@ -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
Numeric/MCMC/Util.hs view
@@ -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 #-}
affine-invariant-ensemble-mcmc.cabal view
@@ -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