module Numeric.MCMC.Flat (
MarkovChain(..), Options(..), Ensemble
, runChain, readInits
) where
import Control.Arrow
import Control.Monad
import Control.Monad.Reader
import Control.Monad.Primitive
import System.Random.MWC
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as U
import Control.Monad.Par (NFData)
import Control.Monad.Par.Scheds.Direct
import Control.Monad.Par.Combinator
import GHC.Float
import System.IO
parMapChunk :: NFData b => Int -> (a -> b) -> [a] -> Par [b]
parMapChunk n f xs = do
xss <- parMap (map f) (chunk n xs)
return (concat xss)
where chunk _ [] = []
chunk m ys = let (as, bs) = splitAt m ys
in as : chunk m bs
data MarkovChain = MarkovChain { ensemble :: Ensemble
, accepts :: !Int }
instance Show MarkovChain where
show config = filter (`notElem` "[]") $ unlines $ map (show . map double2Float) (V.toList (ensemble config))
data Options = Options { _target :: [Double] -> Double
, _size :: !Int
, _csize :: !Int }
type Ensemble = V.Vector [Double]
type ViewsOptions = ReaderT Options
symmetricVariate :: PrimMonad m => Gen (PrimState m) -> m Double
symmetricVariate g = do
z <- uniformR (0 :: Double, 1 :: Double) g
return $! 0.5*(z + 1)^(2 :: Int)
metropolisResult :: [Double] -> [Double]
-> Double -> Double
-> ([Double] -> Double)
-> ([Double], Int)
metropolisResult w0 w1 z zc target =
let val = target proposal target w0 + (fromIntegral (length w0) 1) * log z
proposal = zipWith (+) (map (*z) w0) (map (*(1z)) w1)
in if zc <= min 1 (exp val) then (proposal, 1) else (w0, 0)
executeMoves :: (Functor m, PrimMonad m)
=> Ensemble
-> Ensemble
-> Int
-> Gen (PrimState m)
-> ViewsOptions m (Ensemble, Int)
executeMoves e0 e1 n g = do
Options t _ csize <- ask
zs <- replicateM n (lift $ symmetricVariate g)
zcs <- replicateM n (lift $ uniformR (0 :: Double, 1 :: Double) g)
js <- fmap U.fromList (replicateM n (lift $ uniformR (1:: Int, n) g))
let w0 k = e0 `V.unsafeIndex` (k 1)
w1 k ks = e1 `V.unsafeIndex` ((ks `U.unsafeIndex` (k 1)) 1)
result = runPar $ parMapChunk csize
(\(k, z, zc) -> metropolisResult (w0 k) (w1 k js) z zc t)
(zip3 [1..n] zs zcs)
(newstate, nacc) = (V.fromList . map fst &&& sum . map snd) result
return (newstate, nacc)
metropolisStep :: (Functor m, PrimMonad m)
=> MarkovChain
-> Gen (PrimState m)
-> ViewsOptions m MarkovChain
metropolisStep state g = do
Options _ n _ <- ask
let n0 = truncate (fromIntegral n / (2 :: Double)) :: Int
(e, nacc) = (ensemble &&& accepts) state
(e0, e1) = (V.slice (0 :: Int) n0 &&& V.slice n0 n0) e
result0 <- executeMoves e0 e1 n0 g
result1 <- executeMoves e1 (fst result0) n0 g
return $!
MarkovChain (V.concat $ map fst [result0, result1])
(nacc + snd result0 + snd result1)
runChain :: Options
-> Int
-> Int
-> MarkovChain
-> Gen RealWorld
-> IO MarkovChain
runChain opts nepochs thinEvery initConfig g
| l == 0
= error "runChain: ensemble must contain at least one particle"
| l < (length . V.head) (ensemble initConfig)
= do hPutStrLn stderr $ "runChain: ensemble should be twice as large as "
++ "the target's dimension. Continuing anyway."
go opts nepochs thinEvery initConfig g
| otherwise = go opts nepochs thinEvery initConfig g
where
l = V.length (ensemble initConfig)
go o n t !c g0 | n == 0 = return c
| n `rem` t /= 0 = do
r <- runReaderT (metropolisStep c g0) o
go o (n 1) t r g0
| otherwise = do
r <- runReaderT (metropolisStep c g0) o
print r
go o (n 1) t r g0
readInits :: FilePath -> IO Ensemble
readInits p = fmap (V.fromList . map (map read . words) . lines) (readFile p)