diff --git a/Numeric/MCMC/Metropolis.hs b/Numeric/MCMC/Metropolis.hs --- a/Numeric/MCMC/Metropolis.hs +++ b/Numeric/MCMC/Metropolis.hs @@ -22,6 +22,7 @@ module Numeric.MCMC.Metropolis ( mcmc , chain + , chain' , metropolis -- * Re-exported @@ -61,62 +62,76 @@ metropolis :: (Traversable f, PrimMonad m) => Double + -> Maybe (f Double -> b) -> Transition m (Chain (f Double) b) -metropolis radial = do +metropolis radial tunable = do Chain {..} <- get proposal <- lift (propose radial chainPosition) let proposalScore = lTarget chainTarget proposal acceptProb = whenNaN 0 (exp (min 0 (proposalScore - chainScore))) accept <- lift (MWC.bernoulli acceptProb) - when accept (put (Chain chainTarget proposalScore proposal chainTunables)) + when accept $ do + let tuned = tunable <*> Just proposal + put (Chain chainTarget proposalScore proposal tuned) -- Drive a Markov chain via the Metropolis transition operator. drive :: (Traversable f, PrimMonad m) => Double + -> Maybe (f Double -> b) -> Chain (f Double) b -> Gen (PrimState m) -> Producer (Chain (f Double) b) m c -drive radial = loop where +drive radial tunable = loop where loop state prng = do - next <- lift (MWC.sample (execStateT (metropolis radial) state) prng) + let rvar = execStateT (metropolis radial tunable) state + next <- lift (MWC.sample rvar prng) yield next loop next prng --- | Trace 'n' iterations of a Markov chain and collect the results in a list. --- --- >>> let rosenbrock [x0, x1] = negate (5 *(x1 - x0 ^ 2) ^ 2 + 0.05 * (1 - x0) ^ 2) --- >>> results <- withSystemRandom . asGenIO $ chain 3 1 [0, 0] rosenbrock --- >>> mapM_ print results --- 0.0,0.0 --- 1.4754117657794871e-2,0.5033208261760778 --- 3.8379699517007895e-3,0.24627131099479127 -chain - :: (PrimMonad m, Traversable f) +-- | Return a list of @Chain@ values potentially with tunable values computed +-- from each position. +chain' :: + (PrimMonad m, Traversable f) => Int -> Double -> f Double -> (f Double -> Double) + -> Maybe (f Double -> b) -> Gen (PrimState m) -> m [Chain (f Double) b] -chain n radial position target gen = runEffect $ - drive radial origin gen - >-> collect n +chain' n radial position target tunable gen = + runEffect $ drive radial tunable origin gen >-> collect n where ctarget = Target target Nothing - - origin = Chain { - chainScore = lTarget ctarget position - , chainTunables = Nothing + origin = Chain + { chainScore = lTarget ctarget position + , chainTunables = tunable <*> Just position , chainTarget = ctarget , chainPosition = position } - collect :: Monad m => Int -> Consumer a m [a] - collect size = lowerCodensity $ - replicateM size (lift Pipes.await) + collect size = lowerCodensity $ replicateM size (lift Pipes.await) +-- | Trace 'n' iterations of a Markov chain and collect the results in a list. +-- +-- >>> let rosenbrock [x0, x1] = negate (5 *(x1 - x0 ^ 2) ^ 2 + 0.05 * (1 - x0) ^ 2) +-- >>> results <- withSystemRandom . asGenIO $ chain 3 1 [0, 0] rosenbrock +-- >>> mapM_ print results +-- 0.0,0.0 +-- 1.4754117657794871e-2,0.5033208261760778 +-- 3.8379699517007895e-3,0.24627131099479127 +chain + :: (PrimMonad m, Traversable f) + => Int -- ^ Number of iterations + -> Double -- ^ Step standard deviation + -> f Double -- ^ Starting position + -> (f Double -> Double) -- ^ Log-density (to additive constant) + -> Gen (PrimState m) -- ^ PRNG + -> m [Chain (f Double) b] +chain n radial position target = chain' n radial position target Nothing + -- | Trace 'n' iterations of a Markov chain and stream them to stdout. -- -- >>> let rosenbrock [x0, x1] = negate (5 *(x1 - x0 ^ 2) ^ 2 + 0.05 * (1 - x0) ^ 2) @@ -133,7 +148,7 @@ -> Gen (PrimState m) -> m () mcmc n radial chainPosition target gen = runEffect $ - drive radial Chain {..} gen + drive radial Nothing Chain {..} gen >-> Pipes.take n >-> Pipes.mapM_ (liftIO . print) where diff --git a/mighty-metropolis.cabal b/mighty-metropolis.cabal --- a/mighty-metropolis.cabal +++ b/mighty-metropolis.cabal @@ -1,5 +1,5 @@ name: mighty-metropolis -version: 1.2.0 +version: 2.0.0 synopsis: The Metropolis algorithm. homepage: http://github.com/jtobin/mighty-metropolis license: MIT @@ -8,6 +8,7 @@ maintainer: jared@jtobin.ca category: Numeric build-type: Simple +tested-with: GHC == 8.2.2, GHC == 8.8.3 cabal-version: >= 1.10 description: The classic Metropolis algorithm. @@ -67,7 +68,23 @@ -rtsopts build-depends: base >= 4 && < 6 - , containers >= 0.5 && < 0.6 + , containers >= 0.5 && < 0.7 , mighty-metropolis , mwc-probability >= 1.0.1 +Test-suite tests + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: Spec.hs + default-language: Haskell2010 + ghc-options: + -rtsopts + build-depends: + base >= 4 && < 6 + , containers >= 0.5 && < 0.7 + , foldl + , mighty-metropolis + , mwc-probability >= 1.0.1 + , hspec + , mwc-random + , mcmc-types diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,147 @@ +{-# OPTIONS_GHC -Wall #-} +{-# LANGUAGE RecordWildCards #-} + +import qualified Control.Foldl as L +import Data.Functor.Identity +import Data.Maybe (mapMaybe) +import Data.Sampling.Types +import Numeric.MCMC.Metropolis (chain, chain') +import System.Random.MWC +import Test.Hspec + +withinPercent :: Double -> Double -> Double -> Bool +withinPercent b n a + | b == 0 = a == 0 + | otherwise = d / b < n / 100 + where + d = abs (a - b) + +mean :: [Double] -> Double +mean = L.fold L.mean + +variance :: [Double] -> Double +variance xs = L.fold alg xs where + alg = (/) <$> L.premap csq L.sum <*> L.genericLength - 1 + csq = (** 2.0) . subtract m + m = mean xs + +stdDev :: [Double] -> Double +stdDev = sqrt . variance + +stdErr :: [Double] -> Double +stdErr xs = stdDev xs / sqrt n where + n = fromIntegral (length xs) + +thin :: Int -> [a] -> [a] +thin n xs = case xs of + (h:t) -> h : thin n (drop (pred n) t) + _ -> mempty + +data Params = Params { + pepochs :: Int + , pradial :: Double + , porigin :: Identity Double + , ptunable :: Maybe (Identity Double -> Double) + , pltarget :: Identity Double -> Double + , pthin :: Int + } + +testParams :: Params +testParams = Params { + pepochs = 1000000 + , pradial = 0.2 + , porigin = Identity 1.0 + , ptunable = Just (\(Identity x) -> x ** 3.0) + , pltarget = \(Identity x) -> if x > 0 then negate x else negate 1 / 0 + , pthin = 1000 + } + +vanillaTrace :: IO [Double] +vanillaTrace = do + let Params {..} = testParams + + boxed <- withSystemRandom . asGenIO $ + chain pepochs pradial porigin pltarget + + let positions = fmap (runIdentity . chainPosition) boxed + pure (thin pthin positions) + +tunedTrace :: IO [Double] +tunedTrace = do + let Params {..} = testParams + + boxed <- withSystemRandom . asGenIO $ + chain' pepochs pradial porigin pltarget ptunable + + let positions = mapMaybe chainTunables boxed + pure (thin pthin positions) + +testWithinPercent :: SpecWith () +testWithinPercent = describe "withinPercent" $ + it "works as expected" $ do + 106 `shouldNotSatisfy` withinPercent 100 5 + 105 `shouldNotSatisfy` withinPercent 100 5 + 104 `shouldSatisfy` withinPercent 100 5 + 96 `shouldSatisfy` withinPercent 100 5 + 95 `shouldNotSatisfy` withinPercent 100 5 + 94 `shouldNotSatisfy` withinPercent 100 5 + +testMean :: SpecWith () +testMean = describe "mean" $ + it "works as expected" $ do + mean [1, 2, 3] `shouldSatisfy` withinPercent 2 1e-3 + mean [1..100] `shouldSatisfy` withinPercent 50.5 1e-3 + mean [1..1000000] `shouldSatisfy` withinPercent 500000.5 1e-3 + +testVariance :: SpecWith () +testVariance = describe "variance" $ + it "works as expected" $ do + variance [0, 1] `shouldSatisfy` withinPercent 0.5 1e-3 + variance [1, 1, 1] `shouldSatisfy` withinPercent 0 1e-3 + variance [1..100] `shouldSatisfy` withinPercent 841.66666666 1e-3 + +testStdErr :: SpecWith () +testStdErr = describe "stdErr" $ + it "works as expected" $ do + stdErr [1..100] `shouldSatisfy` withinPercent 2.901149 1e-3 + stdErr [1..1000] `shouldSatisfy` withinPercent 9.133273 1e-3 + +testHelperFunctions :: SpecWith () +testHelperFunctions = describe "helper functions" $ do + testWithinPercent + testMean + testVariance + testStdErr + +testSamples :: [Double] -> SpecWith () +testSamples xs = describe "sampled trace over exp(1)" $ do + let meanStdErr = stdErr xs + varStdErr = stdErr (fmap (\x -> pred x ** 2.0) xs) + + context "within three standard errors" $ do + it "has the expected mean" $ do + mean xs `shouldSatisfy` (< 1 + 3 * meanStdErr) + mean xs `shouldSatisfy` (> 1 - 3 * meanStdErr) + + it "has the expected variance" $ do + variance xs `shouldSatisfy` (< 1 + 3 * varStdErr) + variance xs `shouldSatisfy` (> 1 - 3 * varStdErr) + +testTunables :: [Double] -> SpecWith () +testTunables ts = describe "sampled tunables over exp(1)" $ do + let meanStdErr = stdErr ts + + context "within three standard errors" $ + it "has the expected third moment" $ do + mean ts `shouldSatisfy` (< 6 + 3 * meanStdErr) + mean ts `shouldSatisfy` (> 6 - 3 * meanStdErr) + +main :: IO () +main = do + xs <- vanillaTrace + ts <- tunedTrace + + hspec $ do + testHelperFunctions + testSamples xs + testTunables ts