diff --git a/CHANGES.md b/CHANGES.md --- a/CHANGES.md +++ b/CHANGES.md @@ -1,6 +1,10 @@ Revision history for Haskell package learning-hmm === +## Version 0.3.0.0 +- Add `Learning.IOHMM` which represents a class of input-output HMM +- Delete `Learning.HMM.new` + ## Version 0.2.1.0 - Performance improvements with employing 'hmatrix' - `withEmission` now does iterative re-estimation based on the Viterbi path diff --git a/learning-hmm.cabal b/learning-hmm.cabal --- a/learning-hmm.cabal +++ b/learning-hmm.cabal @@ -1,11 +1,12 @@ name: learning-hmm -version: 0.2.1.0 +version: 0.3.0.0 stability: experimental synopsis: Yet another library for hidden Markov models description: This library provides functions for the maximum likelihood estimation of discrete hidden Markov models. At present, - only Baum-Welch and Viterbi algorithms are implemented. + only Baum-Welch and Viterbi algorithms are implemented for + the plain HMM and the input-output HMM. category: Algorithms, Machine Learning, Statistics author: Mitsuhiro Nakamura @@ -25,10 +26,12 @@ library exposed-modules: Learning.HMM + , Learning.IOHMM other-modules: Data.Random.Distribution.Categorical.Util , Data.Random.Distribution.Simplex , Data.Vector.Generic.Util , Learning.HMM.Internal + , Learning.IOHMM.Internal -- other-extensions: build-depends: base >=4.7 && <4.8 , containers @@ -40,11 +43,3 @@ hs-source-dirs: src default-language: Haskell2010 ghc-options: -Wall - -test-suite doctests - main-is: doctests.hs - type: exitcode-stdio-1.0 - build-depends: base, doctest >= 0.9.11 - hs-source-dirs: tests - default-language: Haskell2010 - ghc-options: -threaded -Wall diff --git a/src/Learning/HMM.hs b/src/Learning/HMM.hs --- a/src/Learning/HMM.hs +++ b/src/Learning/HMM.hs @@ -1,7 +1,6 @@ module Learning.HMM ( HMM (..) , LogLikelihood - , new , init , withEmission , viterbi @@ -12,12 +11,12 @@ import Prelude hiding (init) import Control.Applicative ((<$>)) import Control.Arrow (first) -import Data.List (elemIndex, genericLength) +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, fromWeightedList, normalizeCategoricalPs + fromList, normalizeCategoricalPs ) import Data.Random.Distribution.Categorical.Util () import Data.Random.RVar (RVar) @@ -31,12 +30,27 @@ (!), fromList, fromLists, toList ) import qualified Numeric.LinearAlgebra.HMatrix as H (tr) -import Learning.HMM.Internal +import Learning.HMM.Internal (LogLikelihood) +import qualified Learning.HMM.Internal as I --- | Parameter set of the hidden Markov model. Direct use of the --- constructor is not recommended. Instead, call 'new' or 'init'. -data HMM s o = HMM { states :: [s] -- ^ Hidden states - , outputs :: [o] -- ^ Outputs +-- | Parameter set of the hidden Markov model with discrete emission. +-- The model schema is as follows. +-- +-- @ +-- z_0 -> z_1 -> ... -> z_n +-- | | | +-- v v v +-- x_0 x_1 x_n +-- @ +-- +-- Here, @[z_0, z_1, ..., z_n]@ are hidden states and @[x_0, x_1, ..., x_n]@ +-- are observed outputs. @z_0@ is determined by the 'initialStateDist'. +-- For @i = 1, ..., n@, @z_i@ is determined by the 'transitionDist' +-- conditioned by @z_{i-1}@. +-- For @i = 0, ..., n@, @x_i@ is determined by the 'emissionDist' +-- conditioned by @z_i@. +data HMM s o = HMM { states :: [s] + , outputs :: [o] , initialStateDist :: Categorical Double s -- ^ Categorical distribution of initial state , transitionDist :: s -> Categorical Double s @@ -64,37 +78,11 @@ w = transitionDist hmm phi = emissionDist hmm --- | @new states outputs@ returns a model from the @states@ and @outputs@. --- The 'initialStateDist' and 'emissionDist' are set to be uniform --- distributions. The 'transitionDist' is specified as follows: with --- probability 1/2, move to the same state, otherwise, move to a random --- state (which might be the same state). --- --- >>> new [1, 2 :: Int] ['C', 'D'] --- HMM {states = [1,2], outputs = "CD", initialStateDist = fromList [(0.5,1),(0.5,2)], transitionDist = [(fromList [(0.75,1),(0.25,2)],1),(fromList [(0.25,1),(0.75,2)],2)], emissionDist = [(fromList [(0.5,'C'),(0.5,'D')],1),(fromList [(0.5,'C'),(0.5,'D')],2)]} -new :: (Ord s, Ord o) => [s] -> [o] -> HMM s o -new ss os = HMM { states = ss - , outputs = os - , initialStateDist = pi0 - , transitionDist = w - , emissionDist = phi - } - where - pi0 = C.fromWeightedList [(1, s) | s <- ss] - w s | s `elem` ss = C.fromList [(p s', s') | s' <- ss] - | otherwise = C.fromList [] - where - k = genericLength ss - p s' | s' == s = 1/2 * (1 + 1/k) - | otherwise = 1/2 / k - phi s | s `elem` ss = C.fromWeightedList [(1, o) | o <- os] - | otherwise = C.fromList [] - --- | @init states outputs@ returns a random variable of the model with +-- | @init states outputs@ returns a random variable of models with the -- @states@ and @outputs@, wherein parameters are sampled from uniform -- distributions. init :: (Eq s, Eq o) => [s] -> [o] -> RVar (HMM s o) -init ss os = fromHMM' ss os <$> init' (length ss) (length os) +init ss os = fromInternal ss os <$> I.init (length ss) (length os) -- | @model \`withEmission\` xs@ returns a model in which the -- 'emissionDist' is updated by re-estimations using the observed outputs @@ -102,12 +90,12 @@ -- 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 = fromHMM' ss os $ withEmission' model' xs' +withEmission model xs = fromInternal ss os $ I.withEmission model' xs' where ss = states model os = outputs model os' = V.fromList os - model' = toHMM' model + model' = toInternal model xs' = U.fromList $ fromJust $ mapM (`V.elemIndex` os') xs -- | @viterbi model xs@ performs the Viterbi algorithm using the observed @@ -117,11 +105,11 @@ viterbi model xs = checkModelIn "viterbi" model `seq` checkDataIn "viterbi" model xs `seq` - first toStates $ viterbi' model' xs' + first toStates $ I.viterbi model' xs' where ss' = V.fromList $ states model os' = V.fromList $ outputs model - model' = toHMM' model + model' = toInternal model xs' = U.fromList $ fromJust $ mapM (`V.elemIndex` os') xs toStates = V.toList . V.map (V.unsafeIndex ss') . G.convert @@ -132,16 +120,16 @@ baumWelch model xs = checkModelIn "baumWelch" model `seq` checkDataIn "baumWelch" model xs `seq` - map (first $ fromHMM' ss os) $ baumWelch' model' xs' + map (first $ fromInternal ss os) $ I.baumWelch model' xs' where ss = states model os = outputs model os' = V.fromList os - model' = toHMM' model + model' = toInternal model xs' = U.fromList $ fromJust $ mapM (`V.elemIndex` os') xs -- | @simulate model t@ generates a Markov process of length @t@ using the --- @model@, and returns its state path and observed outputs. +-- @model@, and returns its state path and outputs. simulate :: HMM s o -> Int -> RVar ([s], [o]) simulate model step | step < 1 = return ([], []) @@ -179,35 +167,35 @@ os = outputs hmm err = errorIn fun --- | Convert 'HMM'' to 'HMM'. -fromHMM' :: (Eq s, Eq o) => [s] -> [o] -> HMM' -> HMM s o -fromHMM' 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 - } +-- | 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 + } where - pi0 = initialStateDist' hmm' - w = transitionDist' hmm' - phi = H.tr $ emissionDistT' hmm' + 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 --- | Convert 'HMM' to 'HMM''. The 'initialStateDist'', 'transitionDist'', --- and 'emissionDistT'' are normalized. -toHMM' :: (Eq s, Eq o) => HMM s o -> HMM' -toHMM' hmm = HMM' { nStates' = length ss - , nOutputs' = length os - , initialStateDist' = pi0 - , transitionDist' = w - , emissionDistT' = phi' - } +-- | 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' + } where ss = states hmm os = outputs hmm 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,15 +1,17 @@ module Learning.HMM.Internal ( - HMM' (..) + HMM (..) , LogLikelihood - , init' - , withEmission' - , viterbi' - , baumWelch' - -- , baumWelch1' - -- , forward' - -- , backward' + , init + , withEmission + , viterbi + , baumWelch + -- , baumWelch1 + -- , forward + -- , backward + -- , posterior ) where +import Prelude hiding (init) import Control.Applicative ((<$>)) import Control.DeepSeq (NFData, force, rnf) import Control.Monad (forM_, replicateM) @@ -46,49 +48,49 @@ -- | More efficient data structure of the 'HMM' model. The 'states' and -- 'outputs' in 'HMM' are represented by their indices. The -- 'initialStateDist', 'transitionDist', and 'emissionDist' are --- represented by matrices. The 'emissionDistT'' is a transposed matrix +-- represented by matrices. The 'emissionDistT' is a transposed matrix -- in order to simplify the calculation. -data HMM' = HMM' { nStates' :: Int -- ^ Number of states - , nOutputs' :: Int -- ^ Number of outputs - , initialStateDist' :: H.Vector Double - , transitionDist' :: H.Matrix Double - , emissionDistT' :: H.Matrix Double - } +data HMM = HMM { nStates :: Int -- ^ Number of states + , nOutputs :: Int -- ^ Number of outputs + , initialStateDist :: H.Vector Double + , transitionDist :: H.Matrix Double + , emissionDistT :: H.Matrix Double + } -instance NFData HMM' where - rnf hmm' = rnf k `seq` rnf l `seq` rnf pi0 `seq` rnf w `seq` rnf phi' +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' + k = nStates hmm + l = nOutputs hmm + pi0 = initialStateDist hmm + w = transitionDist hmm + phi' = emissionDistT hmm -init' :: Int -> Int -> RVar HMM' -init' k l = do +init :: Int -> Int -> RVar HMM +init k l = do pi0 <- H.fromList <$> stdSimplex (k-1) w <- H.fromLists <$> replicateM k (stdSimplex (k-1)) phi <- H.fromLists <$> replicateM k (stdSimplex (l-1)) - return HMM' { nStates' = k - , nOutputs' = l - , initialStateDist' = pi0 - , transitionDist' = w - , emissionDistT' = H.tr phi - } + return HMM { nStates = k + , nOutputs = l + , initialStateDist = pi0 + , transitionDist = w + , emissionDistT = H.tr phi + } -withEmission' :: HMM' -> U.Vector Int -> HMM' -withEmission' model xs = model' +withEmission :: HMM -> U.Vector Int -> HMM +withEmission model xs = model' where n = U.length xs - k = nStates' model - l = nOutputs' model + k = nStates model + l = nOutputs model ss = [0..(k-1)] os = [0..(l-1)] - step m = fst $ baumWelch1' (m { emissionDistT' = H.tr phi }) n xs + step m = fst $ baumWelch1 (m { emissionDistT = H.tr phi }) n xs where phi :: H.Matrix Double - phi = let zs = fst $ viterbi' m xs + phi = let zs = fst $ viterbi m xs fs = G.frequencies $ U.zip zs xs hs = H.fromLists $ map (\s -> map (\o -> M.findWithDefault 0 (s, o) fs) os) ss @@ -104,17 +106,17 @@ model' = fst $ head $ dropWhile ((> 1e-9) . snd) $ zip ms' ds -- | Return the Euclidean distance between two models. -euclideanDistance :: HMM' -> HMM' -> Double +euclideanDistance :: HMM -> HMM -> Double euclideanDistance model model' = - sqrt $ (H.sumElements $ (w - w') ** 2) + (H.sumElements $ (phi - phi') ** 2) + sqrt $ H.sumElements ((w - w') ** 2) + H.sumElements ((phi - phi') ** 2) where - w = transitionDist' model - w' = transitionDist' model' - phi = emissionDistT' model - phi' = emissionDistT' model' + w = transitionDist model + w' = transitionDist model' + phi = emissionDistT model + phi' = emissionDistT model' -viterbi' :: HMM' -> U.Vector Int -> (U.Vector Int, LogLikelihood) -viterbi' model xs = (path, logL) +viterbi :: HMM -> U.Vector Int -> (U.Vector Int, LogLikelihood) +viterbi model xs = (path, logL) where n = U.length xs @@ -137,9 +139,9 @@ ps' <- V.unsafeFreeze ps return (ds', ps') where - pi0 = initialStateDist' model - w' = H.toColumns $ transitionDist' model - phi' = emissionDistT' model + pi0 = initialStateDist model + w' = H.toColumns $ transitionDist model + phi' = emissionDistT model deltaE = V.unsafeIndex deltas (n-1) @@ -154,29 +156,29 @@ U.unsafeFreeze ix logL = H.maxElement deltaE -baumWelch' :: HMM' -> U.Vector Int -> [(HMM', LogLikelihood)] -baumWelch' model xs = zip models (tail logLs) +baumWelch :: HMM -> U.Vector Int -> [(HMM, LogLikelihood)] +baumWelch model xs = zip models (tail logLs) where n = U.length xs - step (m, _) = baumWelch1' m n xs + step (m, _) = baumWelch1 m n xs (models, logLs) = unzip $ iterate step (model, undefined) -- | 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 :: HMM -> Int -> U.Vector Int -> (HMM, LogLikelihood) +baumWelch1 model n xs = force (model', logL) where - k = nStates' model - l = nOutputs' model + 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 - betas = backward' model n xs cs + (alphas, cs) = forward model n xs + betas = backward model n xs cs -- Based on the alpha, beta, and scaling values, we calculate the -- posterior distribution, i.e., gamma and xi values. - (gammas, xis) = posterior' model n xs alphas betas cs + (gammas, xis) = posterior model n xs alphas betas cs -- Using the gamma and xi values, we obtain the optimal initial state -- probability vector, transition probability matrix, and emission @@ -191,16 +193,16 @@ in H.fromRows $ map (\o -> ds o / ns) [0..(l-1)] -- We finally obtain the new model and the likelihood for the old model. - model' = model { initialStateDist' = pi0 - , transitionDist' = w - , emissionDistT' = phi' + model' = model { initialStateDist = pi0 + , transitionDist = w + , emissionDistT = phi' } logL = - (U.sum $ U.map log cs) -- | 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 -> Int -> U.Vector Int -> (V.Vector (H.Vector Double), U.Vector Double) +{-# INLINE forward #-} +forward model n xs = runST $ do as <- MV.unsafeNew n cs <- MU.unsafeNew n let x0 = U.unsafeIndex xs 0 @@ -219,15 +221,15 @@ cs' <- U.unsafeFreeze cs return (as', cs') where - k = nStates' model - pi0 = initialStateDist' model - w' = H.tr $ transitionDist' model - phi' = emissionDistT' model + k = nStates model + pi0 = initialStateDist model + w' = H.tr $ transitionDist model + phi' = emissionDistT model -- | 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 -> Int -> U.Vector Int -> U.Vector Double -> V.Vector (H.Vector Double) +{-# INLINE backward #-} +backward model n xs cs = runST $ do bs <- MV.unsafeNew n let bE = H.konst 1 k cE = U.unsafeIndex cs (n-1) @@ -240,19 +242,19 @@ MV.unsafeWrite bs (t-1) (H.konst c' k * b') V.unsafeFreeze bs where - k = nStates' model - w = transitionDist' model - phi' = emissionDistT' model + 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 -> 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) where gammas = V.zipWith3 (\a b c -> a * b / H.konst c k) alphas betas (G.convert cs) xis = V.zipWith3 (\a b x -> H.diag a H.<> w H.<> H.diag (b * (phi' H.! x))) alphas (V.unsafeTail betas) (G.convert $ U.unsafeTail xs) - k = nStates' model - w = transitionDist' model - phi' = emissionDistT' model + k = nStates model + w = transitionDist model + phi' = emissionDistT model diff --git a/src/Learning/IOHMM.hs b/src/Learning/IOHMM.hs new file mode 100644 --- /dev/null +++ b/src/Learning/IOHMM.hs @@ -0,0 +1,240 @@ +module Learning.IOHMM ( + IOHMM (..) + , LogLikelihood + , init + , withEmission + , viterbi + , baumWelch + , 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 + +-- | Parameter set of the input-output hidden Markov model with discrete emission. +-- This 'IOHMM' assumes that the inputs affect only the transition +-- probabilities. The model schema is as follows. +-- +-- @ +-- x_0 x_1 x_n +-- | | +-- v v +-- z_0 -> z_1 -> ... -> z_n +-- | | | +-- v v v +-- y_0 y_1 y_n +-- @ +-- +-- Here, @[x_0, x_1, ..., x_n]@ are given inputs, @[z_0, z_1, ..., z_n]@ +-- are hidden states, and @[y_0, y_1, ..., y_n]@ are observed outputs. +-- @z_0@ is determined by the 'initialStateDist'. +-- For @i = 1, ..., n@, @z_i@ is determined by the 'transitionDist' +-- conditioned by @x_i@ and @z_{i-1}@. +-- For @i = 0, ..., n@, @y_i@ is determined by the 'emissionDist' +-- conditioned by @z_i@. +data IOHMM i s o = IOHMM { inputs :: [i] + , states :: [s] + , outputs :: [o] + , initialStateDist :: Categorical Double s + -- ^ Categorical distribution of initial state + , transitionDist :: i -> s -> Categorical Double s + -- ^ Categorical distribution of next state + -- conditioned by the input and previous state + , emissionDist :: s -> 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, 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 + +-- | @init inputs states outputs@ returns a random variable of models with the +-- @inputs@, @states@, and @outputs@, wherein parameters are sampled from uniform +-- distributions. +init :: (Eq i, Eq s, Eq o) => [i] -> [s] -> [o] -> RVar (IOHMM i s o) +init is ss os = fromInternal is ss os <$> I.init (length is) (length ss) (length os) + +-- | @withEmission model xs ys@ returns a model in which the +-- 'emissionDist' is updated by re-estimations using the inputs @xs@ and +-- outputs @ys@. The 'emissionDist' is set to be normalized histograms +-- each of which is calculated from segumentations of @ys@ based on the +-- Viterbi state path. +-- 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' + 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 + +-- | @viterbi model xs ys@ performs the Viterbi algorithm using the inputs +-- @xs@ and outputs @ys@, and returns the most likely state path and its +-- log likelihood. +-- 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 = + 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 + +-- | @baumWelch model xs ys@ iteratively performs the Baum-Welch algorithm +-- using the inputs @xs@ and outputs @ys@, and returns a list of updated +-- models and their corresponding log likelihoods. +-- 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 = + checkModelIn "baumWelch" model `seq` + checkDataIn "baumWelch" model xs ys `seq` + map (first $ fromInternal is ss os) $ 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 + +-- | @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 + | null xs = return ([], []) + | otherwise = do s0 <- sample $ rvar pi0 + y0 <- sample $ rvar $ phi 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' + ((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 + +-- | 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 + +-- | 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 + } + 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 + +-- | 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' + } + 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] + +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 new file mode 100644 --- /dev/null +++ b/src/Learning/IOHMM/Internal.hs @@ -0,0 +1,269 @@ +module Learning.IOHMM.Internal ( + IOHMM (..) + , LogLikelihood + , init + , withEmission + , viterbi + , baumWelch + -- , baumWelch1 + -- , forward + -- , backward + -- , 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 + ) + +type LogLikelihood = Double + +-- | More efficient data structure of the 'IOHMM' model. The 'inputs', +-- 'states', and 'outputs' in 'IOHMM' are represented by their indices. +-- The 'initialStateDist', 'transitionDist', and 'emissionDist' are +-- represented by matrices. The 'emissionDistT' is a transposed matrix +-- in order to simplify the calculation. +data IOHMM = IOHMM { nInputs :: Int -- ^ Number of inputs + , nStates :: Int -- ^ Number of states + , nOutputs :: Int -- ^ Number of outputs + , initialStateDist :: H.Vector Double + , transitionDist :: V.Vector (H.Matrix Double) + , emissionDistT :: H.Matrix Double + } + +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 + +init :: Int -> Int -> Int -> RVar IOHMM +init m k l = do + pi0 <- H.fromList <$> stdSimplex (k-1) + w <- V.replicateM m (H.fromLists <$> replicateM k (stdSimplex (k-1))) + phi <- H.fromLists <$> replicateM k (stdSimplex (l-1)) + return IOHMM { nInputs = m + , nStates = k + , nOutputs = l + , initialStateDist = pi0 + , transitionDist = w + , emissionDistT = H.tr phi + } + +withEmission :: IOHMM -> U.Vector (Int, Int) -> IOHMM +withEmission model xys = model' + where + n = U.length xys + k = nStates model + l = nOutputs model + ss = [0..(k-1)] + os = [0..(l-1)] + ys = U.map snd xys + + step m = fst $ baumWelch1 (m { emissionDistT = H.tr phi }) n xys + where + phi :: H.Matrix Double + phi = let zs = fst $ viterbi m xys + fs = G.frequencies $ U.zip zs ys + 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) + + ms = iterate step model + ms' = tail ms + ds = zipWith euclideanDistance ms ms' + + model' = fst $ head $ dropWhile ((> 1e-9) . snd) $ zip ms' ds + +-- | Return the Euclidean distance between two models. +euclideanDistance :: IOHMM -> IOHMM -> Double +euclideanDistance model model' = + sqrt $ sum $ H.sumElements ((phi - phi') ** 2) : + map (\i -> H.sumElements ((w i - w' i) ** 2)) is + where + is = [0..(nInputs model - 1)] + w = V.unsafeIndex (transitionDist model) + w' = V.unsafeIndex (transitionDist model') + phi = emissionDistT model + phi' = emissionDistT model' + +viterbi :: IOHMM -> U.Vector (Int, Int) -> (U.Vector Int, LogLikelihood) +viterbi model xys = (path, logL) + where + n = U.length xys + + -- First, we calculate the value function and the state maximizing it + -- for each time. + deltas :: V.Vector (H.Vector Double) + psis :: V.Vector (U.Vector Int) + (deltas, psis) = runST $ do + ds <- MV.unsafeNew n + ps <- MV.unsafeNew n + let (_, y0) = U.unsafeIndex xys 0 + MV.unsafeWrite ds 0 $ log (phi' H.! y0) + log pi0 + 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 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 + + deltaE = V.unsafeIndex deltas (n-1) + + -- The most likely path and corresponding log likelihood are as follows. + path = runST $ do + ix <- MU.unsafeNew n + MU.unsafeWrite ix (n-1) $ H.maxIndex deltaE + forM_ [n-l | l <- [1..(n-1)]] $ \t -> do + i <- MU.unsafeRead ix t + let psi = V.unsafeIndex psis t + MU.unsafeWrite ix (t-1) $ U.unsafeIndex psi i + U.unsafeFreeze ix + logL = H.maxElement deltaE + +baumWelch :: IOHMM -> U.Vector (Int, Int) -> [(IOHMM, LogLikelihood)] +baumWelch model xys = zip models (tail logLs) + where + n = U.length xys + step (m, _) = baumWelch1 m n xys + (models, logLs) = unzip $ iterate step (model, undefined) + +-- | 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) + 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 + -- forward-backward algorithm. + (alphas, cs) = forward model n xys + betas = backward model n xys cs + + -- Based on the alpha, beta, and scaling values, we calculate the + -- posterior distribution, i.e., gamma and xi values. + (gammas, xis) = posterior model n xys alphas betas cs + + -- Using the gamma and xi values, we obtain the optimal initial state + -- probability vector, transition probability matrix, and emission + -- probability matrix. + 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) + 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)] + + -- We finally obtain the new model and the likelihood for the old model. + model' = model { initialStateDist = pi0 + , transitionDist = w + , emissionDistT = phi' + } + logL = - (U.sum $ U.map log cs) + +-- | 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 + as <- MV.unsafeNew n + cs <- MU.unsafeNew n + let (_, y0) = U.unsafeIndex xys 0 + a0 = (phi' H.! y0) * pi0 + c0 = 1 / H.sumElements a0 + MV.unsafeWrite as 0 (H.konst c0 k * 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) + c' = 1 / H.sumElements a' + MV.unsafeWrite as t (H.konst c' k * 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 + +-- | 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 + bs <- MV.unsafeNew n + let bE = H.konst 1 k + cE = U.unsafeIndex cs (n-1) + MV.unsafeWrite bs (n-1) (H.konst cE k * 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) + c' = U.unsafeIndex cs (t-1) + MV.unsafeWrite bs (t-1) (H.konst c' k * b') + V.unsafeFreeze bs + where + k = nStates model + w = V.unsafeIndex (transitionDist model) + phi' = emissionDistT model + +-- | 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) + where + gammas = V.zipWith3 (\a b c -> a * b / H.konst c k) + 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))) + alphas (V.unsafeTail betas) (G.convert $ U.unsafeTail xys) + k = nStates model + w = V.unsafeIndex (transitionDist model) + phi' = emissionDistT model diff --git a/tests/doctests.hs b/tests/doctests.hs deleted file mode 100644 --- a/tests/doctests.hs +++ /dev/null @@ -1,8 +0,0 @@ -module Main (main) where - -import Test.DocTest - -main :: IO () -main = doctest [ "-isrc" - , "src/Learning/HMM.hs" - ]