commit adb4250377e512d2ce0a701845c2805813c900a2
parent b6063d1f7114d212452f60afb4ea5890ef6f7087
Author: Jared Tobin <jared@jtobin.ca>
Date: Wed, 8 Oct 2014 22:30:31 +1300
Version bump.
Diffstat:
7 files changed, 198 insertions(+), 197 deletions(-)
diff --git a/.gitignore b/.gitignore
@@ -13,3 +13,9 @@ build
cabal-dev
site
benchmarks
+.cabal-sandbox
+cabal.sandbox.config
+tmp
+*hp
+*ps
+*aux
diff --git a/Examples/Example.hs b/Examples/Example.hs
@@ -0,0 +1,30 @@
+{-# OPTIONS_GHC -fno-warn-type-defaults #-}
+
+import Data.Foldable (toList)
+import Data.Vector (Vector, fromList)
+import Numeric.AD
+import Numeric.MCMC.Hamiltonian
+import System.Random.MWC
+
+lRosenbrock :: RealFloat a => Vector a -> a
+lRosenbrock xs = negate $ 5 * (x1 - x0 ^ 2) ^ 2 + 0.05 * (1 - x0) ^ 2 where
+ [x0, x1] = toList xs
+
+glRosenbrock :: Vector Double -> Vector Double
+glRosenbrock = grad lRosenbrock
+
+rosenbrock :: Target
+rosenbrock = Target lRosenbrock glRosenbrock
+
+tunables :: Tunables
+tunables = Tunables 0.05 20
+
+options :: Options
+options = Options 5000 (fromList [1, 1])
+
+main :: IO ()
+main = do
+ g <- create
+ trace <- hmc rosenbrock tunables options g
+ mapM_ print trace
+
diff --git a/Examples/code/Rosenbrock_HMC.hs b/Examples/code/Rosenbrock_HMC.hs
@@ -1,47 +0,0 @@
-import Numeric.MCMC.Hamiltonian
-import System.Random.MWC
-import System.Environment
-import System.Exit
-import System.IO
-import Numeric.AD
-import Control.Monad
-
-target :: RealFloat a => [a] -> a
-target [x0, x1] = (-1)*(5*(x1 - x0^2)^2 + 0.05*(1 - x0)^2)
-
-gTarget :: [Double] -> [Double]
-gTarget = grad target
-
-main = do
- args <- getArgs
- when (args == []) $ do
- putStrLn "(hasty-hamiltonian) Rosenbrock density "
- putStrLn "Usage: ./Rosenbrock_HMC <numSteps> <thinEvery> <nDisc> <stepSize> <inits>"
- putStrLn " "
- putStrLn "numSteps : Number of Markov chain iterations to run. "
- putStrLn "thinEvery : Print every nth iteration. "
- putStrLn "nDisc : Number of discretizing steps to take. "
- putStrLn "stepSize : Perturbation scaling parameter. "
- putStrLn "inits : Filepath containing points at which to "
- putStrLn " initialize the chain. "
- exitSuccess
-
- inits <- fmap (map read . words) (readFile (args !! 4)) :: IO [Double]
-
- let nepochs = read (head args) :: Int
- thinEvery = read (args !! 1) :: Int
- nDisc = read (args !! 2) :: Double
- eps = read (args !! 3) :: Double
- params = Options target gTarget nDisc eps
- config = MarkovChain inits 0
-
- g <- create
- results <- runChain params nepochs thinEvery config g
-
- hPutStrLn stderr $
- let nAcc = accepts results
- total = nepochs
- in show nAcc ++ " / " ++ show total ++ " (" ++
- show ((fromIntegral nAcc / fromIntegral total) :: Float) ++
- ") proposals accepted"
-
diff --git a/Examples/code/SPDE_HMC.hs b/Examples/code/SPDE_HMC.hs
@@ -1,55 +0,0 @@
-{-# LANGUAGE BangPatterns #-}
-
-import Numeric.MCMC.Hamiltonian
-import System.Random.MWC
-import System.Environment
-import System.Exit
-import System.IO
-import Control.Monad
-import Numeric.AD
-
-target :: RealFloat a => [a] -> a
-target xs = go 0 0 xs
- where go t0 t1 [] = (- t0 / (2*h)) - (0.5 * h * t1)
- go t0 t1 (u:us:uss) = go (t0 + (us - u)^2) (t1 + v (us + u)) uss
- h = 1 / fromIntegral (length xs)
- v x = (1 - x^2)^2
-{-# INLINE target #-}
-
-gTarget :: [Double] -> [Double]
-gTarget = grad target
-{-# INLINE gTarget #-}
-
-main = do
- args <- getArgs
- when (args == []) $ do
- putStrLn "(hasty-hamiltonian) Stochastic partial differential equation "
- putStrLn "Usage: ./SPDE_HMC <numSteps> <inits> <thinEvery> "
- putStrLn " "
- putStrLn "numSteps : Number of Markov chain iterations to run. "
- putStrLn "thinEvery : Print every nth iteration. "
- putStrLn "nDisc : Number of discretizing steps to take. "
- putStrLn "stepSize : Perturbation scaling parameter. "
- putStrLn "inits : Filepath containing points at which to "
- putStrLn " initialize the chain. "
- exitSuccess
-
- inits <- fmap (map read . words) (readFile (args !! 4)) :: IO [Double]
-
- let nepochs = read (head args) :: Int
- thinEvery = read (args !! 1) :: Int
- nDisc = read (args !! 2) :: Double
- eps = read (args !! 3) :: Double
- params = Options target gTarget nDisc eps
- config = MarkovChain inits 0
-
- g <- create
- results <- runChain params nepochs thinEvery config g
-
- hPutStrLn stderr $
- let nAcc = accepts results
- total = nepochs
- in show nAcc ++ " / " ++ show total ++ " (" ++
- show ((fromIntegral nAcc / fromIntegral total) :: Float) ++
- ") proposals accepted"
-
diff --git a/Numeric/MCMC/Hamiltonian.hs b/Numeric/MCMC/Hamiltonian.hs
@@ -1,88 +1,127 @@
{-# OPTIONS_GHC -Wall #-}
-{-# LANGUAGE BangPatterns #-}
+
+-- | Hamiltonian Monte Carlo. See, ex: Neal (2012)
+-- http://arxiv.org/pdf/1206.1901.pdf.
module Numeric.MCMC.Hamiltonian where
-import Control.Monad.Reader
+import Control.Applicative
+import Control.Monad
import Control.Monad.Primitive
+import Control.Monad.Trans.Class
+import Control.Monad.Trans.State.Strict
+import Data.Vector (Vector)
+import qualified Data.Vector as V
import System.Random.MWC
import System.Random.MWC.Distributions
-import GHC.Float
-
--- | State of the Markov chain. Current parameter values are held in 'theta',
--- while accepts counts the number of proposals accepted.
-data MarkovChain = MarkovChain { theta :: [Double]
- , accepts :: {-# UNPACK #-} !Int }
-
--- | Options for the chain. The target (expected to be a log density), its
--- gradient, and a step size tuning parameter.
-data Options = Options { _target :: [Double] -> Double
- , _gTarget :: [Double] -> [Double]
- , _nLeaps :: {-# UNPACK #-} !Double
- , _eps :: {-# UNPACK #-} !Double }
-
--- | A result with this type has a view of the chain options.
-type ViewsOptions = ReaderT Options
-
--- | Display the current state.
-instance Show MarkovChain where
- show config = filter (`notElem` "[]") $ show (map double2Float (theta config))
-
--- | The 'leapfrog' or Stormer-Verlet discretizer.
-leapfrog :: Monad m
- => [Double]
- -> [Double]
- -> ViewsOptions m ([Double], [Double])
-leapfrog t0 r0 = do
- Options _ gTarget ndisc e <- ask
- let go !t !r n
- | n == 0 = (t, r)
- | otherwise = let rm = zipWith (+) r (map (* (0.5*e)) (gTarget t))
- tt = zipWith (+) t (map (*e) rm)
- rt = zipWith (+) rm (map (* (0.5*e)) (gTarget tt))
- in go tt rt (n - 1)
- return $! go t0 r0 ndisc
-{-# INLINE leapfrog #-}
-
--- | Perform a Metropolis accept/reject step.
-metropolisStep :: PrimMonad m
- => MarkovChain
- -> Gen (PrimState m)
- -> ViewsOptions m MarkovChain
-metropolisStep state g = do
- Options target _ _ _ <- ask
- let t0 = theta state
- nacc = accepts state
- r0 <- replicateM (length t0) (lift $ standard g)
- zc <- lift $ uniformR (0 :: Double, 1 :: Double) g
- (t, r) <- leapfrog t0 r0
-
- let mc = if zc < min 1 (exp arRatio)
- then (t, 1)
- else (t0, 0)
-
- xs <.> ys = sum $ zipWith (*) xs ys
- arRatio = target t + 0.5*(r0 <.> r0)
- - target t0 - 0.5*(r <.> r)
-
- return $! MarkovChain (fst mc) (nacc + snd mc)
-{-# INLINE metropolisStep #-}
-
--- | Diffuse through states.
-runChain :: Options -- Options of the Markov chain.
- -> Int -- Number of epochs to iterate the chain.
- -> Int -- Print every nth iteration.
- -> MarkovChain -- Initial state of the Markov chain.
- -> Gen RealWorld -- MWC PRNG
- -> IO MarkovChain -- End state of the Markov chain, wrapped in IO.
-runChain = go
- where go o n t !c g | n == 0 = return c
- | n `rem` t /= 0 = do
- r <- runReaderT (metropolisStep c g) o
- go o (n - 1) t r g
- | otherwise = do
- r <- runReaderT (metropolisStep c g) o
- print r
- go o (n - 1) t r g
-{-# INLINE runChain #-}
+
+type Parameters = Vector Double
+
+data Target = Target {
+ lTarget :: Parameters -> Double -- ^ log target
+ , glTarget :: Parameters -> Parameters -- ^ gradient of log target
+ }
+
+data Tunables = Tunables {
+ stepSize :: !Double
+ , leapfrogs :: !Int
+ } deriving Show
+
+data Options = Options {
+ epochs :: !Int
+ , initial :: !Parameters
+ } deriving Show
+
+type Particle = (Parameters, Parameters)
+
+type Transition m = StateT Parameters m Parameters
+
+hamiltonian
+ :: PrimMonad m
+ => Target
+ -> Tunables
+ -> Gen (PrimState m)
+ -> Transition m
+hamiltonian target tunables g = do
+ q0 <- get
+ r0 <- V.replicateM (V.length q0) (lift $ standard g)
+ zc <- lift $ uniform g
+ let (q, r) = leapfrogIntegrator target tunables (q0, r0)
+ next = nextState target (q0, q) (r0, r) zc
+ put next
+ return next
+
+leapfrogIntegrator :: Target -> Tunables -> Particle -> Particle
+leapfrogIntegrator target tunables (q0, r0) = go q0 r0 l where
+ l = leapfrogs tunables
+ go q r 0 = (q, r)
+ go q r n =
+ let (q1, r1) = leapfrog target tunables (q, r)
+ in go q1 r1 (pred n)
+
+leapfrog :: Target -> Tunables -> Particle -> Particle
+leapfrog target tunables (q, r) = (qf, rf) where
+ rm = adjustMomentum target tunables (q, r)
+ qf = adjustPosition tunables (rm, q)
+ rf = adjustMomentum target tunables (qf, rm)
+
+-- | Adjust momentum according to a half-leapfrog step.
+adjustMomentum :: Target -> Tunables -> Particle -> Parameters
+adjustMomentum target tunables (q, r) = r .+ ((0.5 * e) .* g q) where
+ e = stepSize tunables
+ g = glTarget target
+
+-- | Adjust position according to a half-leapfrog step.
+adjustPosition :: Tunables -> Particle -> Parameters
+adjustPosition tunables (r, q) = q .+ (e .* r) where
+ e = stepSize tunables
+
+-- | Scalar-vector multiplication.
+(.*) :: Double -> Parameters -> Parameters
+z .* xs = (* z) <$> xs
+
+-- | Scalar-vector subtraction.
+(.-) :: Parameters -> Parameters -> Parameters
+xs .- ys = V.zipWith (-) xs ys
+
+-- | Scalar-vector addition.
+(.+) :: Parameters -> Parameters -> Parameters
+xs .+ ys = V.zipWith (+) xs ys
+
+nextState
+ :: Target
+ -> Particle
+ -> Particle
+ -> Double
+ -> Parameters
+nextState target position momentum z
+ | z < pAccept = snd position
+ | otherwise = fst position
+ where
+ pAccept = acceptProb target position momentum
+
+-- | A target augmented by momentum auxilliary variables.
+auxilliaryTarget :: Target -> Particle -> Double
+auxilliaryTarget target (t, r) = f t - 0.5 * innerProduct r r where
+ f = lTarget target
+
+acceptProb :: Target -> Particle -> Particle -> Double
+acceptProb target (q0, q1) (r0, r1) = exp . min 0 $
+ auxilliaryTarget target (q1, r1) - auxilliaryTarget target (q0, r0)
+
+innerProduct :: Parameters -> Parameters -> Double
+innerProduct xs ys = V.sum $ V.zipWith (*) xs ys
+
+hmc
+ :: PrimMonad m
+ => Target
+ -> Tunables
+ -> Options
+ -> Gen (PrimState m)
+ -> m [Parameters]
+hmc target tunables options g = do
+ let h = hamiltonian target tunables g
+ n = epochs options
+ q = initial options
+ evalStateT (replicateM n h) q
diff --git a/README.md b/README.md
@@ -1,6 +1,16 @@
-# hasty-hamiltonian [](http://travis-ci.org/jtobin/hasty-hamiltonian)
+# hasty-hamiltonian [](http://travis-ci.org/jtobin/hasty-hamiltonian)
-Speedy, gradient-based traversal through parameter space.
+Speedy, gradient-based traversal through parameter space. See the *examples* folder for example usage.
-See the *Examples* folder for example usage.
+Install notes:
+
+ cabal sandbox init
+ cabal install -j --only-dep
+
+To build the example,
+
+ cabal install -j --only-dep --enable-tests
+ cabal build -j hasty-examples
+
+You can then use `./dist/build/hasty-examples/hasty-examples` to see an example trace of a chain wandering over the [Rosenbrock density](http://en.wikipedia.org/wiki/Rosenbrock_function).
diff --git a/hasty-hamiltonian.cabal b/hasty-hamiltonian.cabal
@@ -1,29 +1,47 @@
--- Initial hasty-hamiltonian.cabal generated by cabal init. For further
--- documentation, see http://haskell.org/cabal/users-guide/
-
name: hasty-hamiltonian
-version: 0.1.0.0
+version: 0.2.0.0
synopsis: Speedy traversal through parameter space.
--- description:
homepage: http://jtobin.github.com/hasty-hamiltonian
license: BSD3
license-file: LICENSE
author: Jared Tobin
maintainer: jared@jtobin.ca
--- copyright:
category: Numeric
build-type: Simple
-cabal-version: >=1.8
+cabal-version: >=1.10
Description:
- Speedy gradient-based traversal through parameter spce.
+ Speedy gradient-based traversal through parameter space.
Source-repository head
Type: git
Location: http://github.com/jtobin/hasty-hamiltonian.git
library
- exposed-modules: Numeric.MCMC.Hamiltonian
- -- other-modules:
- build-depends: base ==4.*, mtl ==2.1.*, primitive ==0.4.*, mwc-random ==0.12.*, parallel ==3.2.*
- ghc-options: -Wall
+ default-language: Haskell2010
+ exposed-modules:
+ Numeric.MCMC.Hamiltonian
+ build-depends:
+ base >= 4.7.0.1 && < 5
+ , mtl >= 2.2.1
+ , mwc-random >= 0.13.2.0
+ , primitive >= 0.5.3.0
+ , transformers >= 0.4.1.0
+ , vector >= 0.10.11.0
+ ghc-options:
+ -Wall
+
+Test-Suite hasty-examples
+ type: exitcode-stdio-1.0
+ hs-source-dirs: examples
+ main-is: Example.hs
+ default-language: Haskell2010
+ ghc-options:
+ -threaded -rtsopts -Wall -O2
+ build-depends:
+ ad >= 4.2.1
+ , base >= 4.7.0.1 && < 5
+ , hasty-hamiltonian >= 0.2.0.0
+ , mwc-random >= 0.13.2.0
+ , vector >= 0.10.11.0
+
