diff --git a/CHANGES.md b/CHANGES.md --- a/CHANGES.md +++ b/CHANGES.md @@ -1,6 +1,9 @@ Revision history for Haskell package learning-hmm === +## Version 0.3.1.1 +- Bug fix release + ## Version 0.3.1.0 - Add function `baumWelch'` that performs the Baum-Welch algorithm and returns a result locally maximizing its likelihood. This behaviour is different from diff --git a/LICENSE b/LICENSE --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,6 @@ The MIT License (MIT) -Copyright (c) 2014 Mitsuhiro Nakamura +Copyright (c) 2014-2015 Mitsuhiro Nakamura Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/learning-hmm.cabal b/learning-hmm.cabal --- a/learning-hmm.cabal +++ b/learning-hmm.cabal @@ -1,5 +1,5 @@ name: learning-hmm -version: 0.3.1.0 +version: 0.3.1.1 stability: experimental synopsis: Yet another library for hidden Markov models @@ -11,7 +11,7 @@ author: Mitsuhiro Nakamura maintainer: Mitsuhiro Nakamura -copyright: Copyright (c) 2014 Mitsuhiro Nakamura +copyright: Copyright (c) 2014-2015 Mitsuhiro Nakamura license: MIT license-file: LICENSE homepage: https://github.com/mnacamura/learning-hmm @@ -27,9 +27,9 @@ library exposed-modules: Learning.HMM , Learning.IOHMM - other-modules: Data.Random.Distribution.Categorical.Util + other-modules: Data.Random.Distribution.Extra , Data.Random.Distribution.Simplex - , Data.Vector.Generic.Util + , Data.Vector.Generic.Extra , Learning.HMM.Internal , Learning.IOHMM.Internal -- other-extensions: diff --git a/src/Data/Random/Distribution/Categorical/Util.hs b/src/Data/Random/Distribution/Categorical/Util.hs deleted file mode 100644 --- a/src/Data/Random/Distribution/Categorical/Util.hs +++ /dev/null @@ -1,11 +0,0 @@ -{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses #-} - -module Data.Random.Distribution.Categorical.Util () where - -import Data.Maybe (fromMaybe) -import Data.Random.Distribution (PDF, pdf) -import Data.Random.Distribution.Categorical (Categorical, toList) -import Data.Tuple (swap) - -instance Eq a => PDF (Categorical Double) a where - pdf cat a = fromMaybe 0 (lookup a $ map swap $ toList cat) diff --git a/src/Data/Random/Distribution/Extra.hs b/src/Data/Random/Distribution/Extra.hs new file mode 100644 --- /dev/null +++ b/src/Data/Random/Distribution/Extra.hs @@ -0,0 +1,23 @@ +{-# LANGUAGE + MultiParamTypeClasses, + FlexibleContexts, FlexibleInstances, + UndecidableInstances + #-} + +module Data.Random.Distribution.Extra + ( PMF + , pmf + ) where + +import Data.Maybe ( fromMaybe ) +import Data.Tuple ( swap ) +import Data.Random.Distribution ( Distribution ) +import qualified Data.Random.Distribution.Categorical as C ( Categorical, toList ) + +class Distribution d t => PMF d t where + -- | Probability mass function for discrete distributions + pmf :: d t -> t -> Double + +instance (Real p, Eq a, Distribution (C.Categorical p) a) => PMF (C.Categorical p) a where + pmf d a = realToFrac $ fromMaybe 0 $ lookup a dict + where dict = map swap $ C.toList d diff --git a/src/Data/Random/Distribution/Simplex.hs b/src/Data/Random/Distribution/Simplex.hs --- a/src/Data/Random/Distribution/Simplex.hs +++ b/src/Data/Random/Distribution/Simplex.hs @@ -5,11 +5,11 @@ #-} module Data.Random.Distribution.Simplex - ( StdSimplex(..) - , stdSimplex - , stdSimplexT - , fractionalStdSimplex - ) where + ( StdSimplex (..) + , stdSimplex + , stdSimplexT + , fractionalStdSimplex + ) where import Control.Monad import Data.List @@ -17,32 +17,32 @@ import Data.Random.Distribution import Data.Random.Distribution.Uniform --- |Uniform distribution over a standard simplex. +-- | Uniform distribution over a standard simplex. newtype StdSimplex as = -- | @StdSimplex k@ constructs a standard simplex of dimension @k@ - -- (standard /k/-simplex). - -- An element of the simplex represents a vector variable @as = (a_0, - -- a_1, ..., a_k)@. The elements of @as@ are more than or equal to @0@ - -- and @sum as@ is always equal to @1@. + -- (standard /k/-simplex). + -- An element of the simplex represents a vector variable @as = (a_0, + -- a_1, ..., a_k)@. The elements of @as@ are more than or equal to + -- @0@ and @sum as@ is always equal to @1@. StdSimplex Int deriving (Eq, Show) instance (Ord a, Fractional a, Distribution StdUniform a) => Distribution StdSimplex [a] where - rvar (StdSimplex k) = fractionalStdSimplex k + rvar (StdSimplex k) = fractionalStdSimplex k --- |@stdSimplex k@ returns a random variable being uniformly distributed over --- a standard simplex of dimension @k@. +-- | @stdSimplex k@ returns a random variable being uniformly distributed +-- over a standard simplex of dimension @k@. stdSimplex :: Distribution StdSimplex [a] => Int -> RVar [a] stdSimplex k = rvar (StdSimplex k) stdSimplexT :: Distribution StdSimplex [a] => Int -> RVarT m [a] stdSimplexT k = rvarT (StdSimplex k) --- |An algorithm proposed by Rubinstein & Melamed (1998). --- See, /e.g./, S. Onn, I. Weissman. --- Generating uniform random vectors over a simplex with implications to --- the volume of a certain polytope and to multivariate extremes. --- /Ann Oper Res/ (2011) __189__:331-342. +-- | An algorithm proposed by Rubinstein & Melamed (1998). +-- See, /e.g./, S. Onn, I. Weissman. +-- Generating uniform random vectors over a simplex with implications to +-- the volume of a certain polytope and to multivariate extremes. +-- /Ann Oper Res/ (2011) __189__:331-342. fractionalStdSimplex :: (Ord a, Fractional a, Distribution StdUniform a) => Int -> RVar [a] fractionalStdSimplex k = do us <- replicateM k stdUniform let us' = sort us ++ [1] diff --git a/src/Data/Vector/Generic/Extra.hs b/src/Data/Vector/Generic/Extra.hs new file mode 100644 --- /dev/null +++ b/src/Data/Vector/Generic/Extra.hs @@ -0,0 +1,18 @@ +-- | Extra functions for "Data.Vector" +module Data.Vector.Generic.Extra + ( frequencies + ) where + +import qualified Data.Map.Strict as M ( Map, empty, insertWith ) +import Data.Vector.Generic ( Vector, foldl' ) + +-- $setup +-- >>> :module + Data.Vector + +-- | @frequencies xs@ returns a 'Map' from distinct items in @xs@ to the +-- number of times they appear. +-- +-- >>> frequencies $ fromList "bra bra bar" +-- fromList [(' ',2),('a',3),('b',3),('r',3)] +frequencies :: (Ord a, Vector v a, Num n) => v a -> M.Map a n +frequencies = foldl' (\m k -> M.insertWith (+) k 1 m) M.empty diff --git a/src/Data/Vector/Generic/Util.hs b/src/Data/Vector/Generic/Util.hs deleted file mode 100644 --- a/src/Data/Vector/Generic/Util.hs +++ /dev/null @@ -1,19 +0,0 @@ --- | Miscellaneous utility functions for "Data.Vector" -module Data.Vector.Generic.Util ( - frequencies - ) where - -import Data.Map.Strict (Map) -import qualified Data.Map.Strict as M (empty, insertWith) -import Data.Vector.Generic (Vector, foldl') - --- $setup --- >>> :module + Data.Vector - --- | @frequencies xs@ returns a 'Map' from distinct items in @xs@ to --- the number of times they appear. --- --- >>> frequencies $ fromList "bra bra bar" --- fromList [(' ',2),('a',3),('b',3),('r',3)] -frequencies :: (Ord a, Vector v a, Num n) => v a -> Map a n -frequencies = foldl' (\m k -> M.insertWith (+) k 1 m) M.empty diff --git a/src/Learning/HMM.hs b/src/Learning/HMM.hs --- a/src/Learning/HMM.hs +++ b/src/Learning/HMM.hs @@ -1,5 +1,7 @@ -module Learning.HMM ( - HMM (..) +{-# LANGUAGE RecordWildCards #-} + +module Learning.HMM + ( HMM (..) , LogLikelihood , init , withEmission @@ -9,30 +11,22 @@ , simulate ) where -import Prelude hiding (init) -import Control.Applicative ((<$>)) -import Control.Arrow (first) -import Data.List (elemIndex) -import Data.Maybe (fromJust) -import Data.Random.Distribution (pdf, rvar) -import Data.Random.Distribution.Categorical (Categorical) -import qualified Data.Random.Distribution.Categorical as C ( - fromList, normalizeCategoricalPs - ) -import Data.Random.Distribution.Categorical.Util () -import Data.Random.RVar (RVar) -import Data.Random.Sample (sample) -import qualified Data.Vector as V ( - elemIndex, fromList, map, toList, unsafeIndex - ) -import qualified Data.Vector.Generic as G (convert) -import qualified Data.Vector.Unboxed as U (fromList) -import qualified Numeric.LinearAlgebra.Data as H ( - (!), fromList, fromLists, toList - ) -import qualified Numeric.LinearAlgebra.HMatrix as H (tr) -import Learning.HMM.Internal (LogLikelihood) -import qualified Learning.HMM.Internal as I +import Control.Applicative ( (<$>) ) +import Control.Arrow ( first ) +import Data.List ( elemIndex ) +import Data.Maybe ( fromJust ) +import Data.Random.Distribution ( rvar ) +import qualified Data.Random.Distribution.Categorical as C ( Categorical, fromList, normalizeCategoricalPs ) +import Data.Random.Distribution.Extra ( pmf ) +import Data.Random.RVar ( RVar ) +import qualified Data.Vector as V ( elemIndex, fromList, map, toList, unsafeIndex ) +import qualified Data.Vector.Generic as G ( convert ) +import qualified Data.Vector.Unboxed as U ( fromList ) +import Learning.HMM.Internal ( LogLikelihood ) +import qualified Learning.HMM.Internal as I +import qualified Numeric.LinearAlgebra.Data as H ( (!), fromList, fromLists, toList ) +import qualified Numeric.LinearAlgebra.HMatrix as H ( tr ) +import Prelude hiding ( init ) -- | Parameter set of the hidden Markov model with discrete emission. -- The model schema is as follows. @@ -52,32 +46,23 @@ -- conditioned by @z_i@. data HMM s o = HMM { states :: [s] , outputs :: [o] - , initialStateDist :: Categorical Double s + , initialStateDist :: C.Categorical Double s -- ^ Categorical distribution of initial state - , transitionDist :: s -> Categorical Double s + , transitionDist :: s -> C.Categorical Double s -- ^ Categorical distribution of next state -- conditioned by the previous state - , emissionDist :: s -> Categorical Double o + , emissionDist :: s -> C.Categorical Double o -- ^ Categorical distribution of output conditioned -- by the hidden state } instance (Show s, Show o) => Show (HMM s o) where - show = showHMM - -showHMM :: (Show s, Show o) => HMM s o -> String -showHMM hmm = "HMM {states = " ++ show ss - ++ ", outputs = " ++ show os - ++ ", initialStateDist = " ++ show pi0 - ++ ", transitionDist = " ++ show [(w s, s) | s <- ss] - ++ ", emissionDist = " ++ show [(phi s, s) | s <- ss] - ++ "}" - where - ss = states hmm - os = outputs hmm - pi0 = initialStateDist hmm - w = transitionDist hmm - phi = emissionDist hmm + show HMM {..} = "HMM {states = " ++ show states + ++ ", outputs = " ++ show outputs + ++ ", initialStateDist = " ++ show initialStateDist + ++ ", transitionDist = " ++ show [(transitionDist s, s) | s <- states] + ++ ", emissionDist = " ++ show [(emissionDist s, s) | s <- states] + ++ "}" -- | @init states outputs@ returns a random variable of models with the -- @states@ and @outputs@, wherein parameters are sampled from uniform @@ -91,136 +76,115 @@ -- which is calculated from segumentations of @xs@ based on the Viterbi -- state path. withEmission :: (Eq s, Eq o) => HMM s o -> [o] -> HMM s o -withEmission model xs = fromInternal ss os $ I.withEmission model' xs' +withEmission (model @ HMM {..}) xs = fromInternal states outputs $ I.withEmission model' xs' where - ss = states model - os = outputs model - os' = V.fromList os - model' = toInternal model - xs' = U.fromList $ fromJust $ mapM (`V.elemIndex` os') xs + outputs' = V.fromList outputs + model' = toInternal model + xs' = U.fromList $ fromJust $ mapM (`V.elemIndex` outputs') xs -- | @viterbi model xs@ performs the Viterbi algorithm using the observed -- outputs @xs@, and returns the most likely state path and its log -- likelihood. viterbi :: (Eq s, Eq o) => HMM s o -> [o] -> ([s], LogLikelihood) -viterbi model xs = +viterbi (model @ HMM {..}) xs = checkModelIn "viterbi" model `seq` checkDataIn "viterbi" model xs `seq` first toStates $ I.viterbi model' xs' where - ss' = V.fromList $ states model - os' = V.fromList $ outputs model - model' = toInternal model - xs' = U.fromList $ fromJust $ mapM (`V.elemIndex` os') xs - toStates = V.toList . V.map (V.unsafeIndex ss') . G.convert + states' = V.fromList states + outputs' = V.fromList outputs + model' = toInternal model + xs' = U.fromList $ fromJust $ mapM (`V.elemIndex` outputs') xs + toStates = V.toList . V.map (V.unsafeIndex states') . G.convert -- | @baumWelch model xs@ iteratively performs the Baum-Welch algorithm -- using the observed outputs @xs@, and returns a list of updated models -- and their corresponding log likelihoods. baumWelch :: (Eq s, Eq o) => HMM s o -> [o] -> [(HMM s o, LogLikelihood)] -baumWelch model xs = +baumWelch (model @ HMM {..}) xs = checkModelIn "baumWelch" model `seq` checkDataIn "baumWelch" model xs `seq` - map (first $ fromInternal ss os) $ I.baumWelch model' xs' + map (first $ fromInternal states outputs) $ I.baumWelch model' xs' where - ss = states model - os = outputs model - os' = V.fromList os - model' = toInternal model - xs' = U.fromList $ fromJust $ mapM (`V.elemIndex` os') xs + outputs' = V.fromList outputs + model' = toInternal model + xs' = U.fromList $ fromJust $ mapM (`V.elemIndex` outputs') xs -- | @baumWelch' model xs@ performs the Baum-Welch algorithm using the -- observed outputs @xs@, and returns a model locally maximizing its log -- likelihood. baumWelch' :: (Eq s, Eq o) => HMM s o -> [o] -> (HMM s o, LogLikelihood) -baumWelch' model xs = +baumWelch' (model @ HMM {..}) xs = checkModelIn "baumWelch" model `seq` checkDataIn "baumWelch" model xs `seq` - first (fromInternal ss os) $ I.baumWelch' model' xs' + first (fromInternal states outputs) $ I.baumWelch' model' xs' where - ss = states model - os = outputs model - os' = V.fromList os - model' = toInternal model - xs' = U.fromList $ fromJust $ mapM (`V.elemIndex` os') xs + outputs' = V.fromList outputs + model' = toInternal model + xs' = U.fromList $ fromJust $ mapM (`V.elemIndex` outputs') xs -- | @simulate model t@ generates a Markov process of length @t@ using the -- @model@, and returns its state path and outputs. simulate :: HMM s o -> Int -> RVar ([s], [o]) -simulate model step +simulate HMM {..} step | step < 1 = return ([], []) - | otherwise = do s0 <- sample $ rvar pi0 - x0 <- sample $ rvar $ phi s0 + | otherwise = do s0 <- rvar initialStateDist + x0 <- rvar $ emissionDist s0 unzip . ((s0, x0) :) <$> sim s0 (step - 1) where sim _ 0 = return [] - sim s t = do s' <- sample $ rvar $ w s - x' <- sample $ rvar $ phi s' + sim s t = do s' <- rvar $ transitionDist s + x' <- rvar $ emissionDist s' ((s', x') :) <$> sim s' (t - 1) - pi0 = initialStateDist model - w = transitionDist model - phi = emissionDist model -- | Check if the model is valid in the sense of whether the 'states' and -- 'outputs' are not empty. checkModelIn :: String -> HMM s o -> () -checkModelIn fun hmm - | null ss = err "empty states" - | null os = err "empty outputs" - | otherwise = () - where - ss = states hmm - os = outputs hmm - err = errorIn fun +checkModelIn fun HMM {..} + | null states = errorIn fun "empty states" + | null outputs = errorIn fun "empty outputs" + | otherwise = () -- | Check if all the elements of the observed outputs are contained in the -- 'outputs' of the model. checkDataIn :: Eq o => String -> HMM s o -> [o] -> () -checkDataIn fun hmm xs - | all (`elem` os) xs = () - | otherwise = err "illegal data" - where - os = outputs hmm - err = errorIn fun +checkDataIn fun HMM {..} xs + | all (`elem` outputs) xs = () + | otherwise = errorIn fun "illegal data" -- | Convert internal 'HMM' to 'HMM'. fromInternal :: (Eq s, Eq o) => [s] -> [o] -> I.HMM -> HMM s o -fromInternal ss os hmm' = HMM { states = ss - , outputs = os - , initialStateDist = C.fromList pi0' - , transitionDist = \s -> case elemIndex s ss of - Nothing -> C.fromList [] - Just i -> C.fromList $ w' i - , emissionDist = \s -> case elemIndex s ss of - Nothing -> C.fromList [] - Just i -> C.fromList $ phi' i - } +fromInternal ss os I.HMM {..} = HMM { states = ss + , outputs = os + , initialStateDist = C.fromList pi0' + , transitionDist = \s -> case elemIndex s ss of + Nothing -> C.fromList [] + Just i -> C.fromList $ w' i + , emissionDist = \s -> case elemIndex s ss of + Nothing -> C.fromList [] + Just i -> C.fromList $ phi' i + } where - pi0 = I.initialStateDist hmm' - w = I.transitionDist hmm' - phi = H.tr $ I.emissionDistT hmm' - pi0' = zip (H.toList pi0) ss - w' i = zip (H.toList $ w H.! i) ss - phi' i = zip (H.toList $ phi H.! i) os + pi0' = zip (H.toList initialStateDist) ss + w' i = zip (H.toList $ transitionDist H.! i) ss + phi' i = zip (H.toList $ H.tr emissionDistT H.! i) os -- | Convert 'HMM' to internal 'HMM'. The 'initialStateDist'', -- 'transitionDist'', and 'emissionDistT'' are normalized. toInternal :: (Eq s, Eq o) => HMM s o -> I.HMM -toInternal hmm = I.HMM { I.nStates = length ss - , I.nOutputs = length os - , I.initialStateDist = pi0 - , I.transitionDist = w - , I.emissionDistT = phi' - } +toInternal HMM {..} = I.HMM { I.nStates = length states + , I.nOutputs = length outputs + , I.initialStateDist = pi0 + , I.transitionDist = w + , I.emissionDistT = phi' + } where - ss = states hmm - os = outputs hmm - pi0_ = C.normalizeCategoricalPs $ initialStateDist hmm - w_ = C.normalizeCategoricalPs . transitionDist hmm - phi_ = C.normalizeCategoricalPs . emissionDist hmm - pi0 = H.fromList [pdf pi0_ s | s <- ss] - w = H.fromLists [[pdf (w_ s) s' | s' <- ss] | s <- ss] - phi' = H.fromLists [[pdf (phi_ s) o | s <- ss] | o <- os] + pi0_ = C.normalizeCategoricalPs initialStateDist + w_ = C.normalizeCategoricalPs . transitionDist + phi_ = C.normalizeCategoricalPs . emissionDist + pi0 = H.fromList [pmf pi0_ s | s <- states] + w = H.fromLists [[pmf (w_ s) s' | s' <- states] | s <- states] + phi' = H.fromLists [[pmf (phi_ s) o | s <- states] | o <- outputs] errorIn :: String -> String -> a errorIn fun msg = error $ "Learning.HMM." ++ fun ++ ": " ++ msg diff --git a/src/Learning/HMM/Internal.hs b/src/Learning/HMM/Internal.hs --- a/src/Learning/HMM/Internal.hs +++ b/src/Learning/HMM/Internal.hs @@ -1,5 +1,7 @@ -module Learning.HMM.Internal ( - HMM (..) +{-# LANGUAGE RecordWildCards #-} + +module Learning.HMM.Internal + ( HMM (..) , LogLikelihood , init , withEmission @@ -12,37 +14,22 @@ -- , posterior ) where -import Prelude hiding (init) -import Control.Applicative ((<$>)) -import Control.DeepSeq (NFData, force, rnf) -import Control.Monad (forM_, replicateM) -import Control.Monad.ST (runST) -import qualified Data.Map.Strict as M (findWithDefault) -import Data.Random.RVar (RVar) -import Data.Random.Distribution.Simplex (stdSimplex) -import qualified Data.Vector as V ( - Vector, filter, foldl1', map, unsafeFreeze, unsafeIndex, unsafeTail - , zip, zipWith3 - ) -import qualified Data.Vector.Generic as G (convert) -import qualified Data.Vector.Generic.Util as G (frequencies) -import qualified Data.Vector.Mutable as MV ( - unsafeNew, unsafeRead, unsafeWrite - ) -import qualified Data.Vector.Unboxed as U ( - Vector, fromList, length, map, sum, unsafeFreeze, unsafeIndex - , unsafeTail, zip - ) -import qualified Data.Vector.Unboxed.Mutable as MU ( - unsafeNew, unsafeRead, unsafeWrite - ) -import qualified Numeric.LinearAlgebra.Data as H ( - (!), Matrix, Vector, diag, fromColumns, fromList, fromLists - , fromRows, konst, maxElement, maxIndex, toColumns, tr - ) -import qualified Numeric.LinearAlgebra.HMatrix as H ( - (<>), (#>), sumElements - ) +import Control.Applicative ( (<$>) ) +import Control.DeepSeq ( NFData, force, rnf ) +import Control.Monad ( forM_, replicateM ) +import Control.Monad.ST ( runST ) +import qualified Data.Map.Strict as M ( findWithDefault ) +import Data.Random.Distribution.Simplex ( stdSimplex ) +import Data.Random.RVar ( RVar ) +import qualified Data.Vector as V ( Vector, filter, foldl1', map, unsafeFreeze, unsafeIndex, unsafeTail, zip, zipWith3 ) +import qualified Data.Vector.Generic as G ( convert ) +import qualified Data.Vector.Generic.Extra as G ( frequencies ) +import qualified Data.Vector.Mutable as MV ( unsafeNew, unsafeRead, unsafeWrite ) +import qualified Data.Vector.Unboxed as U ( Vector, fromList, length, map, sum, unsafeFreeze, unsafeIndex, unsafeTail, zip ) +import qualified Data.Vector.Unboxed.Mutable as MU ( unsafeNew, unsafeRead, unsafeWrite ) +import qualified Numeric.LinearAlgebra.Data as H ( (!), Matrix, Vector, diag, fromColumns, fromList, fromLists, fromRows, konst, maxElement, maxIndex, toColumns, tr ) +import qualified Numeric.LinearAlgebra.HMatrix as H ( (<>), (#>), sumElements ) +import Prelude hiding ( init ) type LogLikelihood = Double @@ -59,13 +46,11 @@ } instance NFData HMM where - rnf hmm = rnf k `seq` rnf l `seq` rnf pi0 `seq` rnf w `seq` rnf phi' - where - k = nStates hmm - l = nOutputs hmm - pi0 = initialStateDist hmm - w = transitionDist hmm - phi' = emissionDistT hmm + rnf HMM {..} = rnf nStates `seq` + rnf nOutputs `seq` + rnf initialStateDist `seq` + rnf transitionDist `seq` + rnf emissionDistT init :: Int -> Int -> RVar HMM init k l = do @@ -80,13 +65,11 @@ } withEmission :: HMM -> U.Vector Int -> HMM -withEmission model xs = model' +withEmission (model @ HMM {..}) xs = model' where - n = U.length xs - k = nStates model - l = nOutputs model - ss = [0..(k-1)] - os = [0..(l-1)] + n = U.length xs + ss = [0..(nStates - 1)] + os = [0..(nOutputs - 1)] step m = fst $ baumWelch1 (m { emissionDistT = H.tr phi }) n xs where @@ -96,9 +79,9 @@ hs = H.fromLists $ map (\s -> map (\o -> M.findWithDefault 0 (s, o) fs) os) ss -- hs' is needed to not yield NaN vectors - hs' = hs + H.konst 1e-9 (k, l) - ns = hs' H.#> H.konst 1 k - in hs' / H.fromColumns (replicate l ns) + hs' = hs + H.konst 1e-9 (nStates, nOutputs) + ns = hs' H.#> H.konst 1 nStates + in hs' / H.fromColumns (replicate nOutputs ns) ms = iterate step model ms' = tail ms @@ -117,7 +100,7 @@ phi' = emissionDistT model' viterbi :: HMM -> U.Vector Int -> (U.Vector Int, LogLikelihood) -viterbi model xs = (path, logL) +viterbi HMM {..} xs = (path, logL) where n = U.length xs @@ -129,20 +112,18 @@ ds <- MV.unsafeNew n ps <- MV.unsafeNew n let x0 = U.unsafeIndex xs 0 - MV.unsafeWrite ds 0 $ log (phi' H.! x0) + log pi0 + MV.unsafeWrite ds 0 $ log (emissionDistT H.! x0) + log initialStateDist forM_ [1..(n-1)] $ \t -> do d <- MV.unsafeRead ds (t-1) let x = U.unsafeIndex xs t dws = map (\wj -> d + log wj) w' - MV.unsafeWrite ds t $ log (phi' H.! x) + H.fromList (map H.maxElement dws) + MV.unsafeWrite ds t $ log (emissionDistT H.! x) + H.fromList (map H.maxElement dws) MV.unsafeWrite ps t $ U.fromList (map H.maxIndex dws) ds' <- V.unsafeFreeze ds ps' <- V.unsafeFreeze ps return (ds', ps') where - pi0 = initialStateDist model - w' = H.toColumns $ transitionDist model - phi' = emissionDistT model + w' = H.toColumns transitionDist deltaE = V.unsafeIndex deltas (n-1) @@ -175,11 +156,8 @@ -- | Perform one step of the Baum-Welch algorithm and return the updated -- model and the likelihood of the old model. baumWelch1 :: HMM -> Int -> U.Vector Int -> (HMM, LogLikelihood) -baumWelch1 model n xs = force (model', logL) +baumWelch1 (model @ HMM {..}) n xs = force (model', logL) where - k = nStates model - l = nOutputs model - -- First, we calculate the alpha, beta, and scaling values using the -- forward-backward algorithm. (alphas, cs) = forward model n xs @@ -193,13 +171,13 @@ -- probability vector, transition probability matrix, and emission -- probability matrix. pi0 = V.unsafeIndex gammas 0 - w = let ds = V.foldl1' (+) xis -- denominators - ns = ds H.#> H.konst 1 k -- numerators - in H.diag (H.konst 1 k / ns) H.<> ds + w = let ds = V.foldl1' (+) xis -- denominators + ns = ds H.#> H.konst 1 nStates -- numerators + in H.diag (H.konst 1 nStates / ns) H.<> ds phi' = let gs' o = V.map snd $ V.filter ((== o) . fst) $ V.zip (G.convert xs) gammas ds = V.foldl1' (+) . gs' -- denominators ns = V.foldl1' (+) gammas -- numerators - in H.fromRows $ map (\o -> ds o / ns) [0..(l-1)] + in H.fromRows $ map (\o -> ds o / ns) [0..(nOutputs - 1)] -- We finally obtain the new model and the likelihood for the old model. model' = model { initialStateDist = pi0 @@ -211,59 +189,49 @@ -- | Return alphas and scaling variables. forward :: HMM -> Int -> U.Vector Int -> (V.Vector (H.Vector Double), U.Vector Double) {-# INLINE forward #-} -forward model n xs = runST $ do +forward HMM {..} n xs = runST $ do as <- MV.unsafeNew n cs <- MU.unsafeNew n let x0 = U.unsafeIndex xs 0 - a0 = (phi' H.! x0) * pi0 + a0 = (emissionDistT H.! x0) * initialStateDist c0 = 1 / H.sumElements a0 - MV.unsafeWrite as 0 (H.konst c0 k * a0) + MV.unsafeWrite as 0 (H.konst c0 nStates * a0) MU.unsafeWrite cs 0 c0 forM_ [1..(n-1)] $ \t -> do a <- MV.unsafeRead as (t-1) let x = U.unsafeIndex xs t - a' = (phi' H.! x) * (w' H.#> a) + a' = (emissionDistT H.! x) * (w' H.#> a) c' = 1 / H.sumElements a' - MV.unsafeWrite as t (H.konst c' k * a') + MV.unsafeWrite as t (H.konst c' nStates * a') MU.unsafeWrite cs t c' as' <- V.unsafeFreeze as cs' <- U.unsafeFreeze cs return (as', cs') where - k = nStates model - pi0 = initialStateDist model - w' = H.tr $ transitionDist model - phi' = emissionDistT model + w' = H.tr transitionDist -- | Return betas using scaling variables. backward :: HMM -> Int -> U.Vector Int -> U.Vector Double -> V.Vector (H.Vector Double) {-# INLINE backward #-} -backward model n xs cs = runST $ do +backward HMM {..} n xs cs = runST $ do bs <- MV.unsafeNew n - let bE = H.konst 1 k + let bE = H.konst 1 nStates cE = U.unsafeIndex cs (n-1) - MV.unsafeWrite bs (n-1) (H.konst cE k * bE) + MV.unsafeWrite bs (n-1) (H.konst cE nStates * bE) forM_ [n-l | l <- [1..(n-1)]] $ \t -> do b <- MV.unsafeRead bs t let x = U.unsafeIndex xs t - b' = w H.#> ((phi' H.! x) * b) + b' = transitionDist H.#> ((emissionDistT H.! x) * b) c' = U.unsafeIndex cs (t-1) - MV.unsafeWrite bs (t-1) (H.konst c' k * b') + MV.unsafeWrite bs (t-1) (H.konst c' nStates * b') V.unsafeFreeze bs - where - k = nStates model - w = transitionDist model - phi' = emissionDistT model -- | Return the posterior distribution. posterior :: HMM -> Int -> U.Vector Int -> V.Vector (H.Vector Double) -> V.Vector (H.Vector Double) -> U.Vector Double -> (V.Vector (H.Vector Double), V.Vector (H.Matrix Double)) {-# INLINE posterior #-} -posterior model _ xs alphas betas cs = (gammas, xis) +posterior HMM {..} _ xs alphas betas cs = (gammas, xis) where - gammas = V.zipWith3 (\a b c -> a * b / H.konst c k) + gammas = V.zipWith3 (\a b c -> a * b / H.konst c nStates) alphas betas (G.convert cs) - xis = V.zipWith3 (\a b x -> H.diag a H.<> w H.<> H.diag (b * (phi' H.! x))) + xis = V.zipWith3 (\a b x -> H.diag a H.<> transitionDist H.<> H.diag (b * (emissionDistT H.! x))) alphas (V.unsafeTail betas) (G.convert $ U.unsafeTail xs) - k = nStates model - w = transitionDist model - phi' = emissionDistT model diff --git a/src/Learning/IOHMM.hs b/src/Learning/IOHMM.hs --- a/src/Learning/IOHMM.hs +++ b/src/Learning/IOHMM.hs @@ -1,5 +1,7 @@ -module Learning.IOHMM ( - IOHMM (..) +{-# LANGUAGE RecordWildCards #-} + +module Learning.IOHMM + ( IOHMM (..) , LogLikelihood , init , withEmission @@ -9,30 +11,22 @@ , simulate ) where -import Prelude hiding (init) -import Control.Applicative ((<$>)) -import Control.Arrow (first) -import Data.List (elemIndex) -import Data.Maybe (fromJust) -import Data.Random.Distribution (pdf, rvar) -import Data.Random.Distribution.Categorical (Categorical) -import qualified Data.Random.Distribution.Categorical as C ( - fromList, normalizeCategoricalPs - ) -import Data.Random.Distribution.Categorical.Util () -import Data.Random.RVar (RVar) -import Data.Random.Sample (sample) -import qualified Data.Vector as V ( - elemIndex, fromList, map, toList, unsafeIndex - ) -import qualified Data.Vector.Generic as G (convert) -import qualified Data.Vector.Unboxed as U (fromList, zip) -import qualified Numeric.LinearAlgebra.Data as H ( - (!), fromList, fromLists, toList - ) -import qualified Numeric.LinearAlgebra.HMatrix as H (tr) -import Learning.IOHMM.Internal (LogLikelihood) -import qualified Learning.IOHMM.Internal as I +import Control.Applicative ( (<$>) ) +import Control.Arrow ( first ) +import Data.List ( elemIndex ) +import Data.Maybe ( fromJust ) +import Data.Random.Distribution ( rvar ) +import qualified Data.Random.Distribution.Categorical as C ( Categorical, fromList, normalizeCategoricalPs ) +import Data.Random.Distribution.Extra ( pmf ) +import Data.Random.RVar ( RVar ) +import qualified Data.Vector as V ( elemIndex, fromList, map, toList, unsafeIndex ) +import qualified Data.Vector.Generic as G ( convert ) +import qualified Data.Vector.Unboxed as U ( fromList, zip ) +import qualified Numeric.LinearAlgebra.Data as H ( (!), fromList, fromLists, toList ) +import qualified Numeric.LinearAlgebra.HMatrix as H ( tr ) +import Learning.IOHMM.Internal ( LogLikelihood ) +import qualified Learning.IOHMM.Internal as I +import Prelude hiding ( init ) -- | Parameter set of the input-output hidden Markov model with discrete emission. -- This 'IOHMM' assumes that the inputs affect only the transition @@ -58,34 +52,24 @@ data IOHMM i s o = IOHMM { inputs :: [i] , states :: [s] , outputs :: [o] - , initialStateDist :: Categorical Double s + , initialStateDist :: C.Categorical Double s -- ^ Categorical distribution of initial state - , transitionDist :: i -> s -> Categorical Double s + , transitionDist :: i -> s -> C.Categorical Double s -- ^ Categorical distribution of next state -- conditioned by the input and previous state - , emissionDist :: s -> Categorical Double o + , emissionDist :: s -> C.Categorical Double o -- ^ Categorical distribution of output conditioned -- by the hidden state } instance (Show i, Show s, Show o) => Show (IOHMM i s o) where - show = showIOHMM - -showIOHMM :: (Show i, Show s, Show o) => IOHMM i s o -> String -showIOHMM hmm = "IOHMM {inputs = " ++ show is - ++ ", states = " ++ show ss - ++ ", outputs = " ++ show os - ++ ", initialStateDist = " ++ show pi0 - ++ ", transitionDist = " ++ show [(w i s, (i, s)) | i <- is, s <- ss] - ++ ", emissionDist = " ++ show [(phi s, s) | s <- ss] - ++ "}" - where - is = inputs hmm - ss = states hmm - os = outputs hmm - pi0 = initialStateDist hmm - w = transitionDist hmm - phi = emissionDist hmm + show IOHMM {..} = "IOHMM {inputs = " ++ show inputs + ++ ", states = " ++ show states + ++ ", outputs = " ++ show outputs + ++ ", initialStateDist = " ++ show initialStateDist + ++ ", transitionDist = " ++ show [(transitionDist i s, (i, s)) | i <- inputs, s <- states] + ++ ", emissionDist = " ++ show [(emissionDist s, s) | s <- states] + ++ "}" -- | @init inputs states outputs@ returns a random variable of models with the -- @inputs@, @states@, and @outputs@, wherein parameters are sampled from uniform @@ -101,16 +85,13 @@ -- If the lengths of @xs@ and @ys@ are different, the longer one is cut -- by the length of the shorter one. withEmission :: (Eq i, Eq s, Eq o) => IOHMM i s o -> [i] -> [o] -> IOHMM i s o -withEmission model xs ys = fromInternal is ss os $ I.withEmission model' $ U.zip xs' ys' +withEmission (model @ IOHMM {..}) xs ys = fromInternal inputs states outputs $ I.withEmission model' $ U.zip xs' ys' where - is = inputs model - is' = V.fromList is - ss = states model - os = outputs model - os' = V.fromList os - model' = toInternal model - xs' = U.fromList $ fromJust $ mapM (`V.elemIndex` is') xs - ys' = U.fromList $ fromJust $ mapM (`V.elemIndex` os') ys + inputs' = V.fromList inputs + outputs' = V.fromList outputs + model' = toInternal model + xs' = U.fromList $ fromJust $ mapM (`V.elemIndex` inputs') xs + ys' = U.fromList $ fromJust $ mapM (`V.elemIndex` outputs') ys -- | @viterbi model xs ys@ performs the Viterbi algorithm using the inputs -- @xs@ and outputs @ys@, and returns the most likely state path and its @@ -118,18 +99,18 @@ -- If the lengths of @xs@ and @ys@ are different, the longer one is cut -- by the length of the shorter one. viterbi :: (Eq i, Eq s, Eq o) => IOHMM i s o -> [i] -> [o] -> ([s], LogLikelihood) -viterbi model xs ys = +viterbi (model @ IOHMM {..}) xs ys = checkModelIn "viterbi" model `seq` checkDataIn "viterbi" model xs ys `seq` first toStates $ I.viterbi model' $ U.zip xs' ys' where - is' = V.fromList $ inputs model - ss' = V.fromList $ states model - os' = V.fromList $ outputs model - model' = toInternal model - xs' = U.fromList $ fromJust $ mapM (`V.elemIndex` is') xs - ys' = U.fromList $ fromJust $ mapM (`V.elemIndex` os') ys - toStates = V.toList . V.map (V.unsafeIndex ss') . G.convert + inputs' = V.fromList inputs + states' = V.fromList states + outputs' = V.fromList outputs + model' = toInternal model + xs' = U.fromList $ fromJust $ mapM (`V.elemIndex` inputs') xs + ys' = U.fromList $ fromJust $ mapM (`V.elemIndex` outputs') ys + toStates = V.toList . V.map (V.unsafeIndex states') . G.convert -- | @baumWelch model xs ys@ iteratively performs the Baum-Welch algorithm -- using the inputs @xs@ and outputs @ys@, and returns a list of updated @@ -137,123 +118,99 @@ -- If the lengths of @xs@ and @ys@ are different, the longer one is cut -- by the length of the shorter one. baumWelch :: (Eq i, Eq s, Eq o) => IOHMM i s o -> [i] -> [o] -> [(IOHMM i s o, LogLikelihood)] -baumWelch model xs ys = +baumWelch (model @ IOHMM {..}) xs ys = checkModelIn "baumWelch" model `seq` checkDataIn "baumWelch" model xs ys `seq` - map (first $ fromInternal is ss os) $ I.baumWelch model' $ U.zip xs' ys' + map (first $ fromInternal inputs states outputs) $ I.baumWelch model' $ U.zip xs' ys' where - is = inputs model - is' = V.fromList is - ss = states model - os = outputs model - os' = V.fromList os - model' = toInternal model - xs' = U.fromList $ fromJust $ mapM (`V.elemIndex` is') xs - ys' = U.fromList $ fromJust $ mapM (`V.elemIndex` os') ys + inputs' = V.fromList inputs + outputs' = V.fromList outputs + model' = toInternal model + xs' = U.fromList $ fromJust $ mapM (`V.elemIndex` inputs') xs + ys' = U.fromList $ fromJust $ mapM (`V.elemIndex` outputs') ys -- | @baumWelch' model xs@ performs the Baum-Welch algorithm using the -- inputs @xs@ and outputs @ys@, and returns a model locally maximizing -- its log likelihood. baumWelch' :: (Eq i, Eq s, Eq o) => IOHMM i s o -> [i] -> [o] -> (IOHMM i s o, LogLikelihood) -baumWelch' model xs ys = +baumWelch' (model @ IOHMM {..}) xs ys = checkModelIn "baumWelch" model `seq` checkDataIn "baumWelch" model xs ys `seq` - first (fromInternal is ss os) $ I.baumWelch' model' $ U.zip xs' ys' + first (fromInternal inputs states outputs) $ I.baumWelch' model' $ U.zip xs' ys' where - is = inputs model - is' = V.fromList is - ss = states model - os = outputs model - os' = V.fromList os - model' = toInternal model - xs' = U.fromList $ fromJust $ mapM (`V.elemIndex` is') xs - ys' = U.fromList $ fromJust $ mapM (`V.elemIndex` os') ys + inputs' = V.fromList inputs + outputs' = V.fromList outputs + model' = toInternal model + xs' = U.fromList $ fromJust $ mapM (`V.elemIndex` inputs') xs + ys' = U.fromList $ fromJust $ mapM (`V.elemIndex` outputs') ys -- | @simulate model xs@ generates a Markov process coinciding with the -- inputs @xs@ using the @model@, and returns its state path and observed -- outputs. simulate :: IOHMM i s o -> [i] -> RVar ([s], [o]) -simulate model xs +simulate IOHMM {..} xs | null xs = return ([], []) - | otherwise = do s0 <- sample $ rvar pi0 - y0 <- sample $ rvar $ phi s0 + | otherwise = do s0 <- rvar initialStateDist + y0 <- rvar $ emissionDist s0 unzip . ((s0, y0) :) <$> sim s0 (tail xs) where - sim _ [] = return [] - sim s (x:xs') = do s' <- sample $ rvar $ w x s - y' <- sample $ rvar $ phi s' + sim _ [] = return [] + sim s (x:xs') = do s' <- rvar $ transitionDist x s + y' <- rvar $ emissionDist s' ((s', y') :) <$> sim s' xs' - pi0 = initialStateDist model - w = transitionDist model - phi = emissionDist model -- | Check if the model is valid in the sense of whether the 'states' and -- 'outputs' are not empty. checkModelIn :: String -> IOHMM i s o -> () -checkModelIn fun hmm - | null is = err "empty inputs" - | null ss = err "empty states" - | null os = err "empty outputs" - | otherwise = () - where - is = inputs hmm - ss = states hmm - os = outputs hmm - err = errorIn fun +checkModelIn fun IOHMM {..} + | null inputs = errorIn fun "empty inputs" + | null states = errorIn fun "empty states" + | null outputs = errorIn fun "empty outputs" + | otherwise = () -- | Check if all the elements of the given inputs (outputs) are contained -- in the 'inputs' ('outputs') of the model. checkDataIn :: (Eq i, Eq o) => String -> IOHMM i s o -> [i] -> [o] -> () -checkDataIn fun hmm xs ys - | all (`elem` is) xs && all (`elem` os) ys = () - | otherwise = err "illegal data" - where - is = inputs hmm - os = outputs hmm - err = errorIn fun +checkDataIn fun IOHMM {..} xs ys + | all (`elem` inputs) xs && all (`elem` outputs) ys = () + | otherwise = errorIn fun "illegal data" -- | Convert internal 'IOHMM' to 'IOHMM'. fromInternal :: (Eq i, Eq s, Eq o) => [i] -> [s] -> [o] -> I.IOHMM -> IOHMM i s o -fromInternal is ss os hmm' = IOHMM { inputs = is - , states = ss - , outputs = os - , initialStateDist = C.fromList pi0' - , transitionDist = \i s -> case (elemIndex i is, elemIndex s ss) of - (Nothing, _) -> C.fromList [] - (_, Nothing) -> C.fromList [] - (Just j, Just k) -> C.fromList $ w' j k - , emissionDist = \s -> case elemIndex s ss of - Nothing -> C.fromList [] - Just i -> C.fromList $ phi' i - } +fromInternal is ss os I.IOHMM {..} = IOHMM { inputs = is + , states = ss + , outputs = os + , initialStateDist = C.fromList pi0' + , transitionDist = \i s -> case (elemIndex i is, elemIndex s ss) of + (Nothing, _) -> C.fromList [] + (_, Nothing) -> C.fromList [] + (Just j, Just k) -> C.fromList $ w' j k + , emissionDist = \s -> case elemIndex s ss of + Nothing -> C.fromList [] + Just i -> C.fromList $ phi' i + } where - pi0 = I.initialStateDist hmm' - w = I.transitionDist hmm' - phi = H.tr $ I.emissionDistT hmm' - pi0' = zip (H.toList pi0) ss - w' j k = zip (H.toList $ V.unsafeIndex w j H.! k) ss - phi' i = zip (H.toList $ phi H.! i) os + pi0' = zip (H.toList initialStateDist) ss + w' j k = zip (H.toList $ V.unsafeIndex transitionDist j H.! k) ss + phi' i = zip (H.toList $ H.tr emissionDistT H.! i) os -- | Convert 'IOHMM' to internal 'IOHMM'. The 'initialStateDist'', -- 'transitionDist'', and 'emissionDistT'' are normalized. toInternal :: (Eq i, Eq s, Eq o) => IOHMM i s o -> I.IOHMM -toInternal hmm = I.IOHMM { I.nInputs = length is - , I.nStates = length ss - , I.nOutputs = length os - , I.initialStateDist = pi0 - , I.transitionDist = w - , I.emissionDistT = phi' - } +toInternal IOHMM {..} = I.IOHMM { I.nInputs = length inputs + , I.nStates = length states + , I.nOutputs = length outputs + , I.initialStateDist = pi0 + , I.transitionDist = w + , I.emissionDistT = phi' + } where - is = inputs hmm - ss = states hmm - os = outputs hmm - pi0_ = C.normalizeCategoricalPs $ initialStateDist hmm - w_ i = C.normalizeCategoricalPs . (transitionDist hmm) i - phi_ = C.normalizeCategoricalPs . emissionDist hmm - pi0 = H.fromList [pdf pi0_ s | s <- ss] - w = V.fromList $ map (\i -> H.fromLists [[pdf (w_ i s) s' | s' <- ss] | s <- ss]) is - phi' = H.fromLists [[pdf (phi_ s) o | s <- ss] | o <- os] + pi0_ = C.normalizeCategoricalPs initialStateDist + w_ i = C.normalizeCategoricalPs . transitionDist i + phi_ = C.normalizeCategoricalPs . emissionDist + pi0 = H.fromList [pmf pi0_ s | s <- states] + w = V.fromList $ map (\i -> H.fromLists [[pmf (w_ i s) s' | s' <- states] | s <- states]) inputs + phi' = H.fromLists [[pmf (phi_ s) o | s <- states] | o <- outputs] errorIn :: String -> String -> a errorIn fun msg = error $ "Learning.IOHMM." ++ fun ++ ": " ++ msg diff --git a/src/Learning/IOHMM/Internal.hs b/src/Learning/IOHMM/Internal.hs --- a/src/Learning/IOHMM/Internal.hs +++ b/src/Learning/IOHMM/Internal.hs @@ -1,5 +1,7 @@ -module Learning.IOHMM.Internal ( - IOHMM (..) +{-# LANGUAGE RecordWildCards #-} + +module Learning.IOHMM.Internal + ( IOHMM (..) , LogLikelihood , init , withEmission @@ -12,37 +14,22 @@ -- , posterior ) where -import Prelude hiding (init) -import Control.Applicative ((<$>)) -import Control.DeepSeq (NFData, force, rnf) -import Control.Monad (forM_, replicateM) -import Control.Monad.ST (runST) -import qualified Data.Map.Strict as M (findWithDefault) -import Data.Random.RVar (RVar) -import Data.Random.Distribution.Simplex (stdSimplex) -import qualified Data.Vector as V ( - Vector, filter, foldl1', generate, map, replicateM, unsafeFreeze - , unsafeIndex , unsafeTail , zip, zipWith3 - ) -import qualified Data.Vector.Generic as G (convert) -import qualified Data.Vector.Generic.Util as G (frequencies) -import qualified Data.Vector.Mutable as MV ( - unsafeNew, unsafeRead, unsafeWrite - ) -import qualified Data.Vector.Unboxed as U ( - Vector, fromList, length, map, sum, unsafeFreeze, unsafeIndex - , unsafeTail, unzip, zip - ) -import qualified Data.Vector.Unboxed.Mutable as MU ( - unsafeNew, unsafeRead, unsafeWrite - ) -import qualified Numeric.LinearAlgebra.Data as H ( - (!), Matrix, Vector, diag, fromColumns, fromList, fromLists - , fromRows, konst, maxElement, maxIndex, toColumns, tr - ) -import qualified Numeric.LinearAlgebra.HMatrix as H ( - (<>), (#>), sumElements - ) +import Control.Applicative ( (<$>) ) +import Control.DeepSeq ( NFData, force, rnf ) +import Control.Monad ( forM_, replicateM ) +import Control.Monad.ST ( runST ) +import qualified Data.Map.Strict as M ( findWithDefault ) +import Data.Random.Distribution.Simplex ( stdSimplex ) +import Data.Random.RVar ( RVar ) +import qualified Data.Vector as V ( Vector, filter, foldl1', generate, map, replicateM, unsafeFreeze, unsafeIndex , unsafeTail , zip, zipWith3 ) +import qualified Data.Vector.Generic as G ( convert ) +import qualified Data.Vector.Generic.Extra as G ( frequencies ) +import qualified Data.Vector.Mutable as MV ( unsafeNew, unsafeRead, unsafeWrite ) +import qualified Data.Vector.Unboxed as U ( Vector, fromList, length, map, sum, unsafeFreeze, unsafeIndex, unsafeTail, unzip, zip ) +import qualified Data.Vector.Unboxed.Mutable as MU ( unsafeNew, unsafeRead, unsafeWrite ) +import qualified Numeric.LinearAlgebra.Data as H ( (!), Matrix, Vector, diag, fromColumns, fromList, fromLists, fromRows, konst, maxElement, maxIndex, toColumns, tr ) +import qualified Numeric.LinearAlgebra.HMatrix as H ( (<>), (#>), sumElements ) +import Prelude hiding ( init ) type LogLikelihood = Double @@ -60,14 +47,12 @@ } instance NFData IOHMM where - rnf hmm = rnf m `seq` rnf k `seq` rnf l `seq` rnf pi0 `seq` rnf w `seq` rnf phi' - where - m = nInputs hmm - k = nStates hmm - l = nOutputs hmm - pi0 = initialStateDist hmm - w = transitionDist hmm - phi' = emissionDistT hmm + rnf IOHMM {..} = rnf nInputs `seq` + rnf nStates `seq` + rnf nOutputs `seq` + rnf initialStateDist `seq` + rnf transitionDist `seq` + rnf emissionDistT init :: Int -> Int -> Int -> RVar IOHMM init m k l = do @@ -83,13 +68,11 @@ } withEmission :: IOHMM -> U.Vector (Int, Int) -> IOHMM -withEmission model xys = model' +withEmission (model @ IOHMM {..}) xys = model' where n = U.length xys - k = nStates model - l = nOutputs model - ss = [0..(k-1)] - os = [0..(l-1)] + ss = [0..(nStates - 1)] + os = [0..(nOutputs - 1)] ys = U.map snd xys step m = fst $ baumWelch1 (m { emissionDistT = H.tr phi }) n xys @@ -100,9 +83,9 @@ hs = H.fromLists $ map (\s -> map (\o -> M.findWithDefault 0 (s, o) fs) os) ss -- hs' is needed to not yield NaN vectors - hs' = hs + H.konst 1e-9 (k, l) - ns = hs' H.#> H.konst 1 k - in hs' / H.fromColumns (replicate l ns) + hs' = hs + H.konst 1e-9 (nStates, nOutputs) + ns = hs' H.#> H.konst 1 nStates + in hs' / H.fromColumns (replicate nOutputs ns) ms = iterate step model ms' = tail ms @@ -123,7 +106,7 @@ phi' = emissionDistT model' viterbi :: IOHMM -> U.Vector (Int, Int) -> (U.Vector Int, LogLikelihood) -viterbi model xys = (path, logL) +viterbi IOHMM {..} xys = (path, logL) where n = U.length xys @@ -135,20 +118,18 @@ ds <- MV.unsafeNew n ps <- MV.unsafeNew n let (_, y0) = U.unsafeIndex xys 0 - MV.unsafeWrite ds 0 $ log (phi' H.! y0) + log pi0 + MV.unsafeWrite ds 0 $ log (emissionDistT H.! y0) + log initialStateDist forM_ [1..(n-1)] $ \t -> do d <- MV.unsafeRead ds (t-1) let (x, y) = U.unsafeIndex xys t dws = map (\wj -> d + log wj) (w' x) - MV.unsafeWrite ds t $ log (phi' H.! y) + H.fromList (map H.maxElement dws) + MV.unsafeWrite ds t $ log (emissionDistT H.! y) + H.fromList (map H.maxElement dws) MV.unsafeWrite ps t $ U.fromList (map H.maxIndex dws) ds' <- V.unsafeFreeze ds ps' <- V.unsafeFreeze ps return (ds', ps') where - pi0 = initialStateDist model - w' = H.toColumns . V.unsafeIndex (transitionDist model) - phi' = emissionDistT model + w' = H.toColumns . V.unsafeIndex transitionDist deltaE = V.unsafeIndex deltas (n-1) @@ -181,11 +162,8 @@ -- | Perform one step of the Baum-Welch algorithm and return the updated -- model and the likelihood of the old model. baumWelch1 :: IOHMM -> Int -> U.Vector (Int, Int) -> (IOHMM, LogLikelihood) -baumWelch1 model n xys = force (model', logL) +baumWelch1 (model @ IOHMM {..}) n xys = force (model', logL) where - m = nInputs model - k = nStates model - l = nOutputs model (xs, ys) = U.unzip xys -- First, we calculate the alpha, beta, and scaling values using the @@ -203,12 +181,12 @@ pi0 = V.unsafeIndex gammas 0 w = let xis' i = V.map snd $ V.filter ((== i) . fst) $ V.zip (G.convert $ U.unsafeTail xs) xis ds = V.foldl1' (+) . xis' -- denominators - ns i = ds i H.#> H.konst 1 k -- numerators - in V.map (\i -> H.diag (H.konst 1 k / ns i) H.<> ds i) (V.generate m id) + ns i = ds i H.#> H.konst 1 nStates -- numerators + in V.map (\i -> H.diag (H.konst 1 nStates / ns i) H.<> ds i) (V.generate nInputs id) phi' = let gs' o = V.map snd $ V.filter ((== o) . fst) $ V.zip (G.convert ys) gammas ds = V.foldl1' (+) . gs' -- denominators ns = V.foldl1' (+) gammas -- numerators - in H.fromRows $ map (\o -> ds o / ns) [0..(l-1)] + in H.fromRows $ map (\o -> ds o / ns) [0..(nOutputs - 1)] -- We finally obtain the new model and the likelihood for the old model. model' = model { initialStateDist = pi0 @@ -220,59 +198,52 @@ -- | Return alphas and scaling variables. forward :: IOHMM -> Int -> U.Vector (Int, Int) -> (V.Vector (H.Vector Double), U.Vector Double) {-# INLINE forward #-} -forward model n xys = runST $ do +forward IOHMM {..} n xys = runST $ do as <- MV.unsafeNew n cs <- MU.unsafeNew n let (_, y0) = U.unsafeIndex xys 0 - a0 = (phi' H.! y0) * pi0 + a0 = (emissionDistT H.! y0) * initialStateDist c0 = 1 / H.sumElements a0 - MV.unsafeWrite as 0 (H.konst c0 k * a0) + MV.unsafeWrite as 0 (H.konst c0 nStates * a0) MU.unsafeWrite cs 0 c0 forM_ [1..(n-1)] $ \t -> do a <- MV.unsafeRead as (t-1) let (x, y) = U.unsafeIndex xys t - a' = (phi' H.! y) * (w' x H.#> a) + a' = (emissionDistT H.! y) * (w' x H.#> a) c' = 1 / H.sumElements a' - MV.unsafeWrite as t (H.konst c' k * a') + MV.unsafeWrite as t (H.konst c' nStates * a') MU.unsafeWrite cs t c' as' <- V.unsafeFreeze as cs' <- U.unsafeFreeze cs return (as', cs') where - k = nStates model - pi0 = initialStateDist model - w' = H.tr . V.unsafeIndex (transitionDist model) - phi' = emissionDistT model + w' = H.tr . V.unsafeIndex transitionDist -- | Return betas using scaling variables. backward :: IOHMM -> Int -> U.Vector (Int, Int) -> U.Vector Double -> V.Vector (H.Vector Double) {-# INLINE backward #-} -backward model n xys cs = runST $ do +backward IOHMM {..} n xys cs = runST $ do bs <- MV.unsafeNew n - let bE = H.konst 1 k + let bE = H.konst 1 nStates cE = U.unsafeIndex cs (n-1) - MV.unsafeWrite bs (n-1) (H.konst cE k * bE) + MV.unsafeWrite bs (n-1) (H.konst cE nStates * bE) forM_ [n-l | l <- [1..(n-1)]] $ \t -> do b <- MV.unsafeRead bs t let (x, y) = U.unsafeIndex xys t - b' = w x H.#> ((phi' H.! y) * b) + b' = w x H.#> ((emissionDistT H.! y) * b) c' = U.unsafeIndex cs (t-1) - MV.unsafeWrite bs (t-1) (H.konst c' k * b') + MV.unsafeWrite bs (t-1) (H.konst c' nStates * b') V.unsafeFreeze bs where - k = nStates model - w = V.unsafeIndex (transitionDist model) - phi' = emissionDistT model + w = V.unsafeIndex transitionDist -- | Return the posterior distribution. posterior :: IOHMM -> Int -> U.Vector (Int, Int) -> V.Vector (H.Vector Double) -> V.Vector (H.Vector Double) -> U.Vector Double -> (V.Vector (H.Vector Double), V.Vector (H.Matrix Double)) {-# INLINE posterior #-} -posterior model _ xys alphas betas cs = (gammas, xis) +posterior IOHMM {..} _ xys alphas betas cs = (gammas, xis) where - gammas = V.zipWith3 (\a b c -> a * b / H.konst c k) + gammas = V.zipWith3 (\a b c -> a * b / H.konst c nStates) alphas betas (G.convert cs) - xis = V.zipWith3 (\a b (x, y) -> H.diag a H.<> w x H.<> H.diag (b * (phi' H.! y))) + xis = V.zipWith3 (\a b (x, y) -> H.diag a H.<> w x H.<> H.diag (b * (emissionDistT H.! y))) alphas (V.unsafeTail betas) (G.convert $ U.unsafeTail xys) - k = nStates model - w = V.unsafeIndex (transitionDist model) - phi' = emissionDistT model + w = V.unsafeIndex transitionDist