affine-invariant-ensemble-mcmc (empty) → 0.1.0.0

raw patch · 5 files changed

+299/−0 lines, 5 filesdep +basedep +containersdep +mwc-randomsetup-changed

Dependencies added: base, containers, mwc-random, primitive, split, vector

Files

• LICENSE +29/−0
• Numeric/MCMC/AffineInvariantEnsemble.hs +183/−0
• Numeric/MCMC/Util.hs +59/−0
• Setup.hs +2/−0
• affine-invariant-ensemble-mcmc.cabal +26/−0

@@ -0,0 +1,29 @@+Copyright (C) 2012 Jared Tobin++Redistribution and use in source and binary forms, with or without+modification, are permitted provided that the following conditions are+met:++   1. Redistributions of source code must retain the above copyright+      notice, this list of conditions, and the following disclaimer.++   2. Redistributions in binary form must reproduce the above+      copyright notice, this list of conditions, and the following+      disclaimer in the documentation and/or other materials provided+      with the distribution.++   3. The name of the author may not be used to endorse or promote+      products derived from this software without specific prior+      written permission.++THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR+IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED+WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE+DISCLAIMED. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT,+INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES+(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)+HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,+STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING+IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE+POSSIBILITY OF SUCH DAMAGE.

@@ -0,0 +1,183 @@+{-# LANGUAGE FlexibleInstances, FlexibleContexts #-}++-- | 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. +--+--   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, prune+           ) where++import Numeric.MCMC.Util +import Data.List                        (foldl')+import Data.List.Split                  (splitEvery) +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."++-- | 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.+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]]++instance Show (Trace Double) where+    show (Trace xs) = unlines $ map (unwords . map show) xs++-- 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 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++-- 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)++-- | 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+--   your own initial configuration.+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..] (splitEvery 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)+                        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 possible 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++-- | 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 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.  +--+--   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)+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)++-- | 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+defaultSeed = U.singleton 42 +

@@ -0,0 +1,59 @@+-- | Various internal utilities.+module Numeric.MCMC.Util ++where++import Control.Monad.Primitive (PrimMonad, PrimState)+import System.Random.MWC +import Data.Function           (fix)+import qualified Data.Sequence as Seq+import Data.Foldable           (toList)+import Control.Monad+import Control.Monad.ST+import Data.STRef++-- | 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++-- | Tail-recursive, list-fused 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) ++-- | Convert a list to a pair.+shortListToPair :: [a] -> (a, a)+shortListToPair [x0, x1]  = (x0, x1)+shortListToPair _         = error "shortListToPair - list must have length 2."++-- | Knuth-shuffle a list.  Uses 'Seq' internally.+shuffle :: [a] -> Seed -> [a]+shuffle xs seed = runST $ do+    xsref <- newSTRef $ Seq.fromList xs+    gen   <- restore seed+    let n = length xs+    forM_ [n-1,n-2..1] $ \i -> do+        j <- uniformR (0, i) gen++        xst <- readSTRef xsref+        let tmp = Seq.index xst j++        modifySTRef xsref $ Seq.update j (Seq.index xst i)+        modifySTRef xsref $ Seq.update i tmp++    result <- readSTRef xsref+    return $ toList result++-- | Sample from a list without replacement.+sample :: Int -> [a] -> Seed -> [a]+sample k xs seed = take k $ shuffle xs seed++

@@ -0,0 +1,2 @@+import Distribution.Simple+main = defaultMain

@@ -0,0 +1,26 @@+Name:                affine-invariant-ensemble-mcmc+Homepage:            http://github.com/jtobin/affine-invariant-ensemble-mcmc+Version:             0.1.0.0+Cabal-version:       >=1.8+build-type:          Simple+License:             BSD3+License-file:        LICENSE+Author:              Jared Tobin+Maintainer:          jared@jtobin.ca+Category:            Numerical, Machine Learning, Statistics+Synopsis:            General-purpose sampling +Description:++    A general-purpose sampling routine for badly-scaled 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.*++  ghc-options:         -Wall+  Exposed-modules:     Numeric.MCMC.AffineInvariantEnsemble, Numeric.MCMC.Util+  Extensions:          FlexibleContexts, FlexibleInstances+