diff --git a/BayesStack/Core.hs b/BayesStack/Core.hs new file mode 100644 --- /dev/null +++ b/BayesStack/Core.hs @@ -0,0 +1,7 @@ +module BayesStack.Core( module BayesStack.Core.Types + , module BayesStack.Core.Gibbs + ) where + +import BayesStack.Core.Types +import BayesStack.Core.Gibbs + diff --git a/BayesStack/Core/Gibbs.hs b/BayesStack/Core/Gibbs.hs new file mode 100644 --- /dev/null +++ b/BayesStack/Core/Gibbs.hs @@ -0,0 +1,87 @@ +{-# LANGUAGE TypeFamilies, FlexibleInstances, FlexibleContexts, + ExistentialQuantification, GADTs, CPP #-} + +module BayesStack.Core.Gibbs ( UpdateUnit(..) + , WrappedUpdateUnit(..) + , gibbsUpdate + ) where + +import Control.Monad (replicateM_, when, forever) +import Control.Concurrent +import Control.Concurrent.STM +import GHC.Conc.Sync (labelThread) +import Data.IORef +import Control.DeepSeq +import Data.Random +import Data.Random.Lift +import System.Random.MWC (withSystemRandom) +import Control.Monad.State hiding (lift) + +class UpdateUnit uu where + type ModelState uu + type Setting uu + fetchSetting :: uu -> ModelState uu -> Setting uu + evolveSetting :: ModelState uu -> uu -> RVar (Setting uu) + updateSetting :: uu -> Setting uu -> Setting uu -> ModelState uu -> ModelState uu + +data WrappedUpdateUnit ms = forall uu. (UpdateUnit uu, ModelState uu ~ ms, + NFData (Setting uu), Eq (Setting uu)) + => WrappedUU uu + +updateUnit :: WrappedUpdateUnit ms -> IORef ms -> TBQueue (ms -> ms) -> RVarT IO () +updateUnit (WrappedUU unit) stateRef diffQueue = do + modelState <- lift $ readIORef stateRef + let s = fetchSetting unit modelState + s' <- lift $ evolveSetting modelState unit + (s,s') `deepseq` return () + when (s /= s') $ + lift $ atomically $ writeTBQueue diffQueue (updateSetting unit s s') + +updateWorker :: TQueue (WrappedUpdateUnit ms) -> IORef ms -> TBQueue (ms -> ms) -> RVarT IO () +updateWorker unitsQueue stateRef diffQueue = do + unit <- lift $ atomically $ tryReadTQueue unitsQueue + case unit of + Just unit' -> do updateUnit unit' stateRef diffQueue + updateWorker unitsQueue stateRef diffQueue + Nothing -> return () + +#if __GLASGOW_HASKELL__ < 706 +atomicModifyIORef' = atomicModifyIORef +#endif + +diffWorker :: IORef ms -> TBQueue (ms -> ms) -> Int -> IO () +diffWorker stateRef diffQueue updateBlock = forever $ do + diff <- execStateT (replicateM_ updateBlock $ do + diff <- lift $ atomically $ readTBQueue diffQueue + modify (. diff) + ) id + atomicModifyIORef' stateRef $ \a->(diff a, ()) + +labelMyThread :: String -> IO () +labelMyThread label = myThreadId >>= \id->labelThread id label + +gibbsUpdate :: Int -> ms -> [WrappedUpdateUnit ms] -> IO ms +gibbsUpdate updateBlock modelState units = do + n <- getNumCapabilities + unitsQueue <- atomically $ do q <- newTQueue + mapM_ (writeTQueue q) units + return q + diffQueue <- atomically $ newTBQueue $ 2*updateBlock -- FIXME + stateRef <- newIORef modelState + diffThread <- forkIO $ do labelMyThread "diff worker" + diffWorker stateRef diffQueue updateBlock + + runningWorkers <- atomically $ newTVar (0 :: Int) + done <- atomically $ newEmptyTMVar :: IO (TMVar ()) + replicateM_ n $ forkIO $ withSystemRandom $ \mwc->do + labelMyThread "update worker" + atomically $ modifyTVar' runningWorkers (+1) + runRVarT (updateWorker unitsQueue stateRef diffQueue) mwc + atomically $ do + modifyTVar' runningWorkers (+(-1)) + running <- readTVar runningWorkers + when (running == 0) $ putTMVar done () + + atomically $ takeTMVar done + readIORef stateRef + diff --git a/BayesStack/Core/Types.hs b/BayesStack/Core/Types.hs new file mode 100644 --- /dev/null +++ b/BayesStack/Core/Types.hs @@ -0,0 +1,24 @@ +{-# LANGUAGE TypeFamilies, KindSignatures, ConstraintKinds #-} + +module BayesStack.Core.Types ( Probability + , HasLikelihood(..) + , FullConditionable(..) + ) where + +import GHC.Prim (Constraint) +import Data.Number.LogFloat + +type Probability = LogFloat + +class HasLikelihood p where + type LContext p a :: Constraint + type LContext p a = () + likelihood :: LContext p a => p a -> Probability + prob :: LContext p a => p a -> a -> Probability + +-- | A distribution for which a full conditional factor can be produced +class FullConditionable p where + type FCContext p a :: Constraint + type FCContext p a = () + sampleProb :: FCContext p a => p a -> a -> Double + diff --git a/BayesStack/DirMulti.hs b/BayesStack/DirMulti.hs new file mode 100644 --- /dev/null +++ b/BayesStack/DirMulti.hs @@ -0,0 +1,224 @@ +{-# LANGUAGE TypeFamilies, FlexibleInstances, ConstraintKinds, DeriveGeneric, DefaultSignatures #-} + +module BayesStack.DirMulti ( -- * Dirichlet/multinomial pair + Multinom, dirMulti, symDirMulti, multinom + -- | Do not do record updates with these + , dmTotal, dmAlpha, dmDomain + , setMultinom, SetUnset (..) + , decMultinom, incMultinom + , prettyMultinom + , updatePrior + -- * Parameter estimation + , estimatePrior, reestimatePriors, reestimateSymPriors + -- * Convenience functions + , probabilities, decProbabilities + ) where + +import Data.EnumMap (EnumMap) +import qualified Data.EnumMap as EM + +import Data.Sequence (Seq) +import qualified Data.Sequence as SQ + +import qualified Data.Foldable +import Data.Foldable (toList, Foldable, foldMap) +import Data.Function (on) + +import Text.PrettyPrint +import Text.Printf + +import GHC.Generics (Generic) +import Data.Serialize +import Data.Serialize.EnumMap () +import Data.Serialize.LogFloat () + +import BayesStack.Core +import BayesStack.Dirichlet + +import Data.Number.LogFloat hiding (realToFrac, isNaN, isInfinite) +import Numeric.Digamma +import Math.Gamma hiding (p) + +-- | Make error handling a bit easier +checkNaN :: RealFloat a => String -> a -> a +checkNaN loc x | isNaN x = error $ "BayesStack.DirMulti."++loc++": Not a number" +checkNaN loc x | isInfinite x = error $ "BayesStack.DirMulti."++loc++": Infinity" +checkNaN _ x = x + +maybeInc, maybeDec :: Maybe Int -> Maybe Int +maybeInc Nothing = Just 1 +maybeInc (Just n) = Just (n+1) +maybeDec Nothing = error "Can't decrement zero count" +maybeDec (Just 1) = Nothing +maybeDec (Just n) = Just (n-1) + +{-# INLINEABLE decMultinom #-} +{-# INLINEABLE incMultinom #-} +decMultinom, incMultinom :: (Ord a, Enum a) => a -> Multinom a -> Multinom a +decMultinom k dm = dm { dmCounts = EM.alter maybeDec k $ dmCounts dm + , dmTotal = dmTotal dm - 1 } +incMultinom k dm = dm { dmCounts = EM.alter maybeInc k $ dmCounts dm + , dmTotal = dmTotal dm + 1 } + +data SetUnset = Set | Unset + +setMultinom :: (Enum a, Ord a) => SetUnset -> a -> Multinom a -> Multinom a +setMultinom Set s = incMultinom s +setMultinom Unset s = decMultinom s + +-- | 'Multinom a' represents multinomial distribution over domain 'a'. +-- Optionally, this can include a collapsed Dirichlet prior. +-- 'Multinom alpha count total' is a multinomial with Dirichlet prior +-- with symmetric parameter 'alpha', ... +data Multinom a = DirMulti { dmAlpha :: Alpha a + , dmCounts :: EnumMap a Int + , dmTotal :: !Int + , dmDomain :: Seq a + } + | Multinom { dmProbs :: !(EnumMap a Double) + , dmCounts :: !(EnumMap a Int) + , dmTotal :: !Int + , dmDomain :: !(Seq a) + } + deriving (Show, Eq, Generic) +instance (Enum a, Serialize a) => Serialize (Multinom a) + +-- | 'symMultinomFromPrecision d p' is a symmetric Dirichlet/multinomial over a +-- domain 'd' with precision 'p' +symDirMultiFromPrecision :: Enum a => [a] -> DirPrecision -> Multinom a +symDirMultiFromPrecision domain prec = symDirMulti (0.5*prec) domain + +-- | 'dirMultiFromMeanPrecision m p' is an asymmetric Dirichlet/multinomial +-- over a domain 'd' with mean 'm' and precision 'p' +dirMultiFromPrecision :: Enum a => DirMean a -> DirPrecision -> Multinom a +dirMultiFromPrecision m p = dirMultiFromAlpha $ meanPrecisionToAlpha m p + +-- | Create a symmetric Dirichlet/multinomial +symDirMulti :: Enum a => Double -> [a] -> Multinom a +symDirMulti alpha domain = dirMultiFromAlpha $ symAlpha domain alpha + +-- | A multinomial without a prior +multinom :: Enum a => [(a,Double)] -> Multinom a +multinom probs = Multinom { dmProbs = EM.fromList probs + , dmCounts = EM.empty + , dmTotal = 0 + , dmDomain = SQ.fromList $ map fst probs + } + +-- | Create an asymmetric Dirichlet/multinomial from items and alphas +dirMulti :: Enum a => [(a,Double)] -> Multinom a +dirMulti domain = dirMultiFromAlpha $ asymAlpha $ EM.fromList domain + +-- | Create a Dirichlet/multinomial with a given prior +dirMultiFromAlpha :: Enum a => Alpha a -> Multinom a +dirMultiFromAlpha alpha = DirMulti { dmAlpha = alpha + , dmCounts = EM.empty + , dmTotal = 0 + , dmDomain = alphaDomain alpha + } + +dmGetCounts :: Enum a => Multinom a -> a -> Int +dmGetCounts dm k = + EM.findWithDefault 0 k (dmCounts dm) + +instance HasLikelihood Multinom where + type LContext Multinom a = (Ord a, Enum a) + likelihood dm@(Multinom {}) = + product $ map (\(k,n)->(realToFrac $ dmProbs dm EM.! k)^n) $ EM.assocs $ dmCounts dm + likelihood dm = + let alpha = dmAlpha dm + f k = logToLogFloat $ checkNaN "likelihood(factor)" + $ lnGamma (realToFrac (dmGetCounts dm k) + alpha `alphaOf` k) + in 1 / alphaNormalizer alpha + * product (map f $ toList $ dmDomain dm) + / logToLogFloat (checkNaN "likelihood" $ lnGamma $ realToFrac (dmTotal dm) + sumAlpha alpha) + {-# INLINEABLE likelihood #-} + + prob dm@(Multinom {}) k = realToFrac $ dmProbs dm EM.! k + prob dm k = + let alpha = dmAlpha dm + f k = logToLogFloat $ checkNaN "prob(factor)" + $ lnGamma (realToFrac (dmGetCounts dm k) + alpha `alphaOf` k) + in 1 / alphaNormalizer alpha + * f k + / logToLogFloat (checkNaN "prob" $ lnGamma $ realToFrac (dmTotal dm) + sumAlpha alpha) + {-# INLINEABLE prob #-} + +instance FullConditionable Multinom where + type FCContext Multinom a = (Ord a, Enum a) + sampleProb (Multinom {dmProbs=prob}) k = prob EM.! k + sampleProb dm@(DirMulti {dmAlpha=a}) k = + let alpha = a `alphaOf` k + n = realToFrac $ dmGetCounts dm k + total = realToFrac $ dmTotal dm + in (n + alpha) / (total + sumAlpha a) + {-# INLINEABLE sampleProb #-} + +{-# INLINEABLE probabilities #-} +probabilities :: (Ord a, Enum a) => Multinom a -> Seq (Double, a) +probabilities dm = fmap (\a->(sampleProb dm a, a)) $ dmDomain dm -- FIXME + +-- | Probabilities sorted decreasingly +decProbabilities :: (Ord a, Enum a) => Multinom a -> Seq (Double, a) +decProbabilities = SQ.sortBy (flip (compare `on` fst)) . probabilities + +prettyMultinom :: (Ord a, Enum a) => Int -> (a -> String) -> Multinom a -> Doc +prettyMultinom _ _ (Multinom {}) = error "TODO: prettyMultinom" +prettyMultinom n showA dm@(DirMulti {}) = + text "DirMulti" <+> parens (text "alpha=" <> prettyAlpha showA (dmAlpha dm)) + $$ nest 5 (fsep $ punctuate comma + $ map (\(p,a)->text (showA a) <> parens (text $ printf "%1.2e" p)) + $ take n $ Data.Foldable.toList $ decProbabilities dm) + +-- | Update the prior of a Dirichlet/multinomial +updatePrior :: (Alpha a -> Alpha a) -> Multinom a -> Multinom a +updatePrior _ (Multinom {}) = error "TODO: updatePrior" +updatePrior f dm = dm {dmAlpha=f $ dmAlpha dm} + +-- | Relative tolerance in precision for prior estimation +estimationTol = 1e-8 + +reestimatePriors :: (Foldable f, Functor f, Enum a) => f (Multinom a) -> f (Multinom a) +reestimatePriors dms = + let usableDms = filter (\dm->dmTotal dm > 5) $ toList dms + alpha = case () of + _ | length usableDms <= 3 -> id + otherwise -> const $ estimatePrior estimationTol usableDms + in fmap (updatePrior alpha) dms + +reestimateSymPriors :: (Foldable f, Functor f, Enum a) => f (Multinom a) -> f (Multinom a) +reestimateSymPriors dms = + let usableDms = filter (\dm->dmTotal dm > 5) $ toList dms + alpha = case () of + _ | length usableDms <= 3 -> id + otherwise -> const $ symmetrizeAlpha $ estimatePrior estimationTol usableDms + in fmap (updatePrior alpha) dms + +-- | Estimate the prior alpha from a set of Dirichlet/multinomials +estimatePrior' :: (Enum a) => [Multinom a] -> Alpha a -> Alpha a +estimatePrior' dms alpha = + let domain = toList $ dmDomain $ head dms + f k = let num = sum $ map (\i->digamma (realToFrac (dmGetCounts i k) + alphaOf alpha k) + - digamma (alphaOf alpha k) + ) + $ filter (\i->dmGetCounts i k > 0) dms + total i = realToFrac $ sum $ map (\k->dmGetCounts i k) domain + sumAlpha = sum $ map (alphaOf alpha) domain + denom = sum $ map (\i->digamma (total i + sumAlpha) - digamma sumAlpha) dms + in case () of + _ | isNaN num -> error $ "BayesStack.DirMulti.estimatePrior': num = NaN: "++show (map (\i->(digamma (realToFrac (dmGetCounts i k) + alphaOf alpha k), digamma (alphaOf alpha k))) dms) + _ | denom == 0 -> error "BayesStack.DirMulti.estimatePrior': denom=0" + _ | isInfinite num -> error "BayesStack.DirMulti.estimatePrior': num is infinity " + _ | isNaN (alphaOf alpha k * num / denom) -> error $ "NaN"++show (num, denom) + otherwise -> alphaOf alpha k * num / denom + in asymAlpha $ foldMap (\k->EM.singleton k (f k)) domain + +estimatePrior :: (Enum a) => Double -> [Multinom a] -> Alpha a +estimatePrior tol dms = iter $ dmAlpha $ head dms + where iter alpha = let alpha' = estimatePrior' dms alpha + (_, prec) = alphaToMeanPrecision alpha + (_, prec') = alphaToMeanPrecision alpha' + in if abs ((prec' - prec) / prec) > tol + then iter alpha' + else alpha' + diff --git a/BayesStack/Dirichlet.hs b/BayesStack/Dirichlet.hs new file mode 100644 --- /dev/null +++ b/BayesStack/Dirichlet.hs @@ -0,0 +1,144 @@ +{-# LANGUAGE DeriveGeneric #-} + +module BayesStack.Dirichlet ( -- * Dirichlet parameter + Alpha + , symAlpha, asymAlpha + , alphaDomain, alphaNormalizer, sumAlpha + , DirMean, DirPrecision + , alphaOf, setAlphaOf, setSymAlpha + , alphaToMeanPrecision, meanPrecisionToAlpha + , symmetrizeAlpha + , prettyAlpha + ) where + +import Data.Foldable (toList, Foldable, fold) + +import Data.EnumMap (EnumMap) +import qualified Data.EnumMap as EM + +import Data.Sequence (Seq) +import qualified Data.Sequence as SQ + +import Data.Number.LogFloat hiding (realToFrac, isNaN, isInfinite) +import Math.Gamma + +import Text.Printf +import Text.PrettyPrint + +import Data.Serialize +import Data.Serialize.EnumMap () +import Data.Serialize.LogFloat () +import GHC.Generics (Generic) + +-- | Make error handling a bit easier +checkNaN :: RealFloat a => String -> a -> a +checkNaN loc x | isNaN x = error $ "BayesStack.Dirichlet."++loc++": Not a number" +checkNaN loc x | isInfinite x = error $ "BayesStack.Dirichlet."++loc++": Infinity" +checkNaN _ x = x + +-- | A Dirichlet prior +data Alpha a = SymAlpha { aDomain :: Seq a + , aAlpha :: !Double + , aNorm :: LogFloat + } + | Alpha { aAlphas :: EnumMap a Double + , aSumAlphas :: !Double + , aNorm :: LogFloat + } + deriving (Show, Eq, Generic) +instance (Enum a, Serialize a) => Serialize (Alpha a) + +type DirMean a = EnumMap a Double +type DirPrecision = Double + +symAlpha :: Enum a => [a] -> Double -> Alpha a +symAlpha domain _ | null domain = error "Dirichlet over null domain is undefined" +symAlpha domain alpha = SymAlpha { aDomain = SQ.fromList domain + , aAlpha = alpha + , aNorm = alphaNorm $ symAlpha domain alpha + } + +-- | Construct an asymmetric Alpha +asymAlpha :: Enum a => EnumMap a Double -> Alpha a +asymAlpha alphas | EM.null alphas = error "Dirichlet over null domain is undefined" +asymAlpha alphas = Alpha { aAlphas = alphas + , aSumAlphas = sum $ EM.elems alphas + , aNorm = alphaNorm $ asymAlpha alphas + } + +setSymAlpha :: Enum a => Double -> Alpha a -> Alpha a +setSymAlpha alpha a = let b = (symmetrizeAlpha a) { aAlpha = alpha + , aNorm = alphaNorm b + } + in b + +-- | Compute the normalizer of the likelihood involving alphas, +-- (product_k gamma(alpha_k)) / gamma(sum_k alpha_k) +alphaNorm :: Enum a => Alpha a -> LogFloat +alphaNorm alpha = normNum / normDenom + where dim = realToFrac $ SQ.length $ aDomain alpha + normNum = case alpha of + Alpha {} -> product $ map (\a->logToLogFloat $ checkNaN ("alphaNorm.normNum(asym) alpha="++show a) $ lnGamma a) + $ EM.elems $ aAlphas alpha + SymAlpha {} -> logToLogFloat $ checkNaN "alphaNorm.normNum(sym)" $ dim * lnGamma (aAlpha alpha) + normDenom = logToLogFloat $ checkNaN "alphaNorm.normDenom" $ lnGamma $ sumAlpha alpha + +-- | 'alphaDomain a' is the domain of prior 'a' +alphaDomain :: Enum a => Alpha a -> Seq a +alphaDomain (SymAlpha {aDomain=d}) = d +alphaDomain (Alpha {aAlphas=a}) = SQ.fromList $ EM.keys a + +alphaNormalizer :: Enum a => Alpha a -> LogFloat +alphaNormalizer = aNorm + +-- | 'alphaOf alpha k' is the value of element 'k' in prior 'alpha' +alphaOf :: Enum a => Alpha a -> a -> Double +alphaOf (SymAlpha {aAlpha=alpha}) = const alpha +alphaOf (Alpha {aAlphas=alphas}) = (alphas EM.!) + +-- | 'sumAlpha alpha' is the sum of all alphas +sumAlpha :: Enum a => Alpha a -> Double +sumAlpha (SymAlpha {aDomain=domain, aAlpha=alpha}) = realToFrac (SQ.length domain) * alpha +sumAlpha (Alpha {aSumAlphas=sum}) = sum + +-- | Set a particular alpha element +setAlphaOf :: Enum a => a -> Double -> Alpha a -> Alpha a +setAlphaOf k a alpha@(SymAlpha {}) = setAlphaOf k a $ asymmetrizeAlpha alpha +setAlphaOf k a (Alpha {aAlphas=alphas}) = asymAlpha $ EM.insert k a alphas + +-- | 'alphaToMeanPrecision a' is the mean/precision representation of the prior 'a' +alphaToMeanPrecision :: Enum a => Alpha a -> (DirMean a, DirPrecision) +alphaToMeanPrecision (SymAlpha {aDomain=dom, aAlpha=alpha}) = + let prec = realToFrac (SQ.length dom) * alpha + in (EM.fromList $ map (\a->(a, alpha/prec)) $ toList dom, prec) +alphaToMeanPrecision (Alpha {aAlphas=alphas, aSumAlphas=prec}) = + (fmap (/prec) alphas, prec) + +-- | 'meanPrecisionToAlpha m p' is a prior with mean 'm' and precision 'p' +meanPrecisionToAlpha :: Enum a => DirMean a -> DirPrecision -> Alpha a +meanPrecisionToAlpha mean prec = asymAlpha $ fmap (*prec) mean + +-- | Symmetrize a Dirichlet prior (such that mean=0) +symmetrizeAlpha :: Enum a => Alpha a -> Alpha a +symmetrizeAlpha alpha@(SymAlpha {}) = alpha +symmetrizeAlpha alpha@(Alpha {}) = + SymAlpha { aDomain = alphaDomain alpha + , aAlpha = sumAlpha alpha / realToFrac (EM.size $ aAlphas alpha) + , aNorm = alphaNorm $ symmetrizeAlpha alpha + } + +-- | Turn a symmetric alpha into an asymmetric alpha. For internal use. +asymmetrizeAlpha :: Enum a => Alpha a -> Alpha a +asymmetrizeAlpha (SymAlpha {aDomain=domain, aAlpha=alpha}) = + asymAlpha $ fold $ fmap (\k->EM.singleton k alpha) domain +asymmetrizeAlpha alpha@(Alpha {}) = alpha + +-- | Pretty-print a Dirichlet prior +prettyAlpha :: Enum a => (a -> String) -> Alpha a -> Doc +prettyAlpha showA (SymAlpha {aAlpha=alpha}) = text "Symmetric" <+> double alpha +prettyAlpha showA (Alpha {aAlphas=alphas}) = + text "Assymmetric" + <+> fsep (punctuate comma + $ map (\(a,alpha)->text (showA a) <> parens (text $ printf "%1.2e" alpha)) + $ take 100 $ EM.toList $ alphas) + diff --git a/BayesStack/TupleEnum.hs b/BayesStack/TupleEnum.hs new file mode 100644 --- /dev/null +++ b/BayesStack/TupleEnum.hs @@ -0,0 +1,9 @@ +module BayesStack.TupleEnum where + +-- This is probably a bad idea +-- Perhaps some TH magic to automatically derive this would make it more tolerable +instance (Enum a, Enum b) => Enum (a,b) where + fromEnum (a,b) = 2^32 * fromEnum a + fromEnum b + toEnum n = let (na, nb) = n `quotRem` (2^32) + in (toEnum na, toEnum nb) + diff --git a/BayesStack/UniqueKey.hs b/BayesStack/UniqueKey.hs new file mode 100644 --- /dev/null +++ b/BayesStack/UniqueKey.hs @@ -0,0 +1,84 @@ +{-# LANGUAGE GeneralizedNewtypeDeriving, CPP #-} + +module BayesStack.UniqueKey ( getUniqueKey + , getValueMap, getKeyMap + , mapTraversable + , UniqueKey, UniqueKeyT + , runUniqueKey, runUniqueKeyT + , runUniqueKey', runUniqueKeyT' + ) where + +import Prelude hiding (mapM) +import Control.Applicative (Applicative, (<$>)) +import Data.Traversable (Traversable, mapM) +import Data.Tuple +import Data.Functor.Identity + +import Control.Monad.Trans +import Control.Monad.State.Strict hiding (mapM) + +#if __GLASGOW_HASKELL__ >= 706 +import Data.Map.Strict (Map) +import qualified Data.Map.Strict as M +#else +import Data.Map (Map) +import qualified Data.Map as M +#endif + +-- | 'UniqueKey val key' is a monad for a calculation of a mapping unique keys +-- 'key' onto values 'val' +type UniqueKey val key = UniqueKeyT val key Identity +newtype UniqueKeyT val key m a = UniqueKeyT (StateT ([key], Map val key) m a) + deriving (Monad, Applicative, Functor, MonadTrans) + +-- | Get map of unique keys to values +getKeyMap :: (Monad m, Applicative m, Ord key, Ord val) => UniqueKeyT val key m (Map key val) +getKeyMap = M.fromList . map swap . M.toList <$> getValueMap + +-- | Get map of values to unique keys +getValueMap :: (Monad m, Applicative m, Ord key, Ord val) => UniqueKeyT val key m (Map val key) +getValueMap = snd <$> UniqueKeyT get + +popUniqueKey :: Monad m => UniqueKeyT val key m key +popUniqueKey = do + (keys, a) <- UniqueKeyT get + case keys of + key:rest -> UniqueKeyT (put $! (rest, a)) >> return key + [] -> error "Ran out of unique keys" + +-- | Find the unique key for value 'val' or 'Nothing' if the value is unknown +findUniqueKey :: (Monad m, Applicative m, Ord key, Ord val) => val -> UniqueKeyT val key m (Maybe key) +findUniqueKey value = M.lookup value <$> getValueMap + +getUniqueKey :: (Monad m, Applicative m, Ord key, Ord val) => val -> UniqueKeyT val key m key +getUniqueKey x = do + key <- findUniqueKey x + case key of + Just k -> return k + Nothing -> do k <- popUniqueKey + UniqueKeyT $ modify $ \(keys, keyMap)->(keys, M.insert x k keyMap) + return k + +runUniqueKey :: (Ord key) => [key] -> UniqueKey val key a -> a +runUniqueKey keys = runIdentity . runUniqueKeyT keys + +runUniqueKeyT :: (Monad m, Ord key) => [key] -> UniqueKeyT val key m a -> m a +runUniqueKeyT keys (UniqueKeyT a) = evalStateT a (keys, M.empty) + +-- | Run a `UniqueKeyT`, returning the result and the associated key map +runUniqueKeyT' :: (Monad m, Applicative m, Ord key, Ord val) => [key] -> UniqueKeyT val key m a -> m (a, Map key val) +runUniqueKeyT' keys action = + runUniqueKeyT keys $ do result <- action + keyMap <- getKeyMap + return (result, keyMap) + +-- | Run a `UniqueKey`, returning the result and the associated key map +runUniqueKey' :: (Ord key, Ord val) => [key] -> UniqueKey val key a -> (a, Map key val) +runUniqueKey' keys action = + runUniqueKey keys $ do result <- action + keyMap <- getKeyMap + return (result, keyMap) + +mapTraversable :: (Traversable t, Ord key, Ord val) => [key] -> t val -> (t key, Map key val) +mapTraversable keys xs = runUniqueKey' keys $ mapM getUniqueKey xs + diff --git a/Data/Random/Sequence.hs b/Data/Random/Sequence.hs new file mode 100644 --- /dev/null +++ b/Data/Random/Sequence.hs @@ -0,0 +1,11 @@ +module Data.Random.Sequence (randomElementT) where + +import Data.Random +import Data.Sequence as SQ + +randomElementT :: Seq a -> RVarT m a +randomElementT xs | SQ.null xs = error "randomElementT: empty seq!" +randomElementT xs = do + n <- uniformT 0 (SQ.length xs - 1) + return (xs `index` n) + diff --git a/Data/Sequence/Chunk.hs b/Data/Sequence/Chunk.hs new file mode 100644 --- /dev/null +++ b/Data/Sequence/Chunk.hs @@ -0,0 +1,11 @@ +module Data.Sequence.Chunk (chunk) where + +import Data.Sequence as SQ + +-- | 'chunk n xs' splits 'xs' into 'n' chunks +chunk :: Int -> Seq a -> Seq (Seq a) +chunk n xs = let m = ceiling $ realToFrac (SQ.length xs) / realToFrac n + f xs | SQ.null xs = SQ.empty + f xs = SQ.take m xs <| (f $ SQ.drop m xs) + in f xs + diff --git a/Data/Serialize/EnumMap.hs b/Data/Serialize/EnumMap.hs new file mode 100644 --- /dev/null +++ b/Data/Serialize/EnumMap.hs @@ -0,0 +1,10 @@ +module Data.Serialize.EnumMap where + +import Data.Serialize +import Data.EnumMap + +instance (Enum k, Serialize k, Serialize v) => Serialize (EnumMap k v) where + get = do a <- get + return $ fromList a + put = put . toList + diff --git a/Data/Serialize/LogFloat.hs b/Data/Serialize/LogFloat.hs new file mode 100644 --- /dev/null +++ b/Data/Serialize/LogFloat.hs @@ -0,0 +1,9 @@ +module Data.Serialize.LogFloat where + +import Data.Serialize +import Data.Number.LogFloat + +instance Serialize LogFloat where + put = put . (logFromLogFloat :: LogFloat -> Double) + get = (logToLogFloat :: Double -> LogFloat) `fmap` get + diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright (c)2012, Ben Gamari, Laura Dietz + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * 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. + + * Neither the name of Ben Gamari nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"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 COPYRIGHT +OWNER OR CONTRIBUTORS 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/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/bayes-stack.cabal b/bayes-stack.cabal new file mode 100644 --- /dev/null +++ b/bayes-stack.cabal @@ -0,0 +1,52 @@ +Name: bayes-stack + +Version: 0.2.0.1 +Synopsis: Framework for inferring generative probabilistic models + with Gibbs sampling +Description: bayes-stack is a framework for inference on generative + probabilistic models. The framework uses Gibbs sampling, + although is suitable for other iterative update methods. +homepage: https://github.com/bgamari/bayes-stack +License: BSD3 +License-file: LICENSE +Author: Ben Gamari +Maintainer: bgamari.foss@gmail.com +copyright: Copyright (c) 2012 Ben Gamari +Category: Math + +Build-type: Simple +Cabal-version: >=1.6 + +source-repository head + type: git + location: https://github.com/bgamari/bayes-stack.git + +Library + Exposed-modules: BayesStack.Core, BayesStack.Core.Types, BayesStack.Core.Gibbs, + BayesStack.DirMulti, BayesStack.Dirichlet, + BayesStack.UniqueKey, + BayesStack.TupleEnum, + Data.Serialize.EnumMap, + Data.Serialize.LogFloat, + Data.Random.Sequence, Data.Sequence.Chunk + + Build-depends: base >=4 && <5, + stm, + transformers, + mtl, + deepseq, + random-source, + random-fu, + rvar, + containers, + enummapset, + ghc-prim, + vector, + mwc-random, + pretty, + cereal, + logfloat, + digamma, + gamma, + statistics +