diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -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. diff --git a/Numeric/MCMC/AffineInvariantEnsemble.hs b/Numeric/MCMC/AffineInvariantEnsemble.hs new file mode 100644 --- /dev/null +++ b/Numeric/MCMC/AffineInvariantEnsemble.hs @@ -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 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 . +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 + diff --git a/Numeric/MCMC/Util.hs b/Numeric/MCMC/Util.hs new file mode 100644 --- /dev/null +++ b/Numeric/MCMC/Util.hs @@ -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 + + diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/affine-invariant-ensemble-mcmc.cabal b/affine-invariant-ensemble-mcmc.cabal new file mode 100644 --- /dev/null +++ b/affine-invariant-ensemble-mcmc.cabal @@ -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 +