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/Flat.hs b/Numeric/MCMC/Flat.hs new file mode 100644 --- /dev/null +++ b/Numeric/MCMC/Flat.hs @@ -0,0 +1,157 @@ +{-# OPTIONS_GHC -Wall #-} +{-# LANGUAGE BangPatterns #-} + +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 + +-- | Parallel map with a specified granularity. +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 + +-- | State of the Markov chain. Current ensemble position is held in 'theta', +-- while 'accepts' counts the number of proposals accepted. +data MarkovChain = MarkovChain { ensemble :: Ensemble + , accepts :: {-# UNPACK #-} !Int } + +-- | Display the current state. This will be very slow and should be replaced. +instance Show MarkovChain where + show config = filter (`notElem` "[]") $ unlines $ map (show . map double2Float) (V.toList (ensemble config)) + +-- | Options for the chain. The target (expected to be a log density), as +-- well as the size of the ensemble. The size should be an even number. Also +-- holds the specified parallel granularity as 'csize'. +data Options = Options { _target :: [Double] -> Double + , _size :: {-# UNPACK #-} !Int + , _csize :: {-# UNPACK #-} !Int } + +-- | An ensemble of particles. +type Ensemble = V.Vector [Double] + +-- | A result with this type has a view of the chain's options. +type ViewsOptions = ReaderT Options + +-- | Generate a random value from a distribution having the property that +-- g(1/z) = zg(z). +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) +{-# INLINE symmetricVariate #-} + +-- | The result of a single-particle Metropolis accept/reject step. This +-- compares a particle's state to a perturbation made by an affine +-- transformation based on a complementary particle. Non-monadic to +-- more easily be used in the Par monad. +metropolisResult :: [Double] -> [Double] -- Target and alternate particles + -> Double -> Double -- z ~ g(z) and zc ~ rand + -> ([Double] -> Double) -- Target function + -> ([Double], Int) -- Result and accept counter +metropolisResult w0 w1 z zc target = + let val = target proposal - target w0 + (fromIntegral (length w0) - 1) * log z + proposal = zipWith (+) (map (*z) w0) (map (*(1-z)) w1) + in if zc <= min 1 (exp val) then (proposal, 1) else (w0, 0) +{-# INLINE metropolisResult #-} + +-- | Execute Metropolis steps on the particles of a sub-ensemble by +-- perturbing them with affine transformations based on particles +-- in a complementary ensemble, in parallel. +executeMoves :: (Functor m, PrimMonad m) + => Ensemble -- Target sub-ensemble + -> Ensemble -- Complementary sub-ensemble + -> Int -- Size of the sub-ensembles + -> Gen (PrimState m) -- MWC PRNG + -> ViewsOptions m (Ensemble, Int) -- Updated ensemble and # of accepts +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) +{-# INLINE executeMoves #-} + +-- | Perform a Metropolis accept/reject step on the ensemble by +-- perturbing each element and accepting/rejecting the perturbation in +-- parallel. +metropolisStep :: (Functor m, PrimMonad m) + => MarkovChain -- State of the Markov chain + -> Gen (PrimState m) -- MWC PRNG + -> ViewsOptions m MarkovChain -- Updated sub-ensemble +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 + + -- Update each sub-ensemble + 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) +{-# 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 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 +{-# INLINE runChain #-} + +-- | A convenience function to read and parse ensemble inits from disk. +-- Assumes a text file with one particle per line, where each particle +-- element is separated by whitespace. +readInits :: FilePath -> IO Ensemble +readInits p = fmap (V.fromList . map (map read . words) . lines) (readFile p) +{-# INLINE readInits #-} + 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/flat-mcmc.cabal b/flat-mcmc.cabal new file mode 100644 --- /dev/null +++ b/flat-mcmc.cabal @@ -0,0 +1,31 @@ +-- Initial flat_mcmc.cabal generated by cabal init. For further +-- documentation, see http://haskell.org/cabal/users-guide/ + +name: flat-mcmc +version: 0.1.0.0 +synopsis: Painless general-purpose sampling. +-- description: +homepage: http://jtobin.github.com/flat-mcmc +-- license: +license: BSD3 +license-file: LICENSE +author: Jared Tobin +maintainer: jared@jtobin.ca +-- copyright: +category: Numeric, Machine Learning, Statistics +build-type: Simple +cabal-version: >=1.8 +Description: + + Painless general-purpose sampling. + +Source-repository head + Type: git + Location: http://github.com/jtobin/flat-mcmc.git + +library + exposed-modules: Numeric.MCMC.Flat + -- other-modules: + build-depends: base ==4.*, mtl ==2.1.*, primitive ==0.4.*, mwc-random ==0.12.*, vector ==0.9.*, monad-par ==0.3.*, monad-par-extras ==0.3.* + ghc-options: -Wall +