# Hidden Markov Model

# 2.1 Implementation

The Wikipedia page gives an imperative implementation in pseudo-code and a detailed worked example:

https://en.wikipedia.org/wiki/Viterbi_algorithm

I ported it to Haskell, with the sparse transition matrix represented as a Map of Maps storing log-probabilities (otherwise probabilities get too small and underflow to 0).

I needed to transpose the matrix from my constructed “prefix -> suffix” order, as the algorithm needs to find the most probable path ending at each given suffix.

I also returned the computed log-probability, a value of -Infinity means there was no match at all and the returned path will be garbage (occurs with probability 0).

# 2.2 Partial Observations

One use case I have in mind is generating sequences that are fixed at particular points, but are otherwise unconstrained.

A partial / unconstrained / uninformative observation has an equal probability connection to every internal state.

It seems exact probabilities are not strictly necessary, so giving them all the same arbitrary weight works ok (I used probability 1).

# 2.3 List Of Matches

Haskell’s lazy evaluation makes it straightforward to extend the algorithm to return all matching sequences sorted by probability (most likely first), from which you can take as many as needed.

Instead of storing the most likely previous state at each state at each time step, store a list of (log-probability, path prefix) pairs, with the invariant that they are sorted in descending order. I stored the paths in reverse order to benefit from lazy sharing of tails, and correct that as a post-processing step.

Then the key operation is merging a list of sorted lists, which can be done by folding over the outer list with a pairwise merge.

# 3.1 Fever

See the Wikipedia page linked above.

My experiment with two observed dizzy days spaced apart showed that 10 days would persist a fever but at 11 days apart it would collapse into two single fever days with a healthy period between.

# 3.2 Text

The internal state is 3 sequential characters of text (trigrams).

The observation of a state is the middle character of the three with probability 1, all other characters having probability 0.

The transition matrix is constructed from a long input text sequence (I tested with the Adventures of Sherlock Holmes by Arthur Conan Doyle from Project Gutenberg).

# 4.1 Viterbi

module Viterbi where

import Prelude hiding (init)

import Data.List (maximumBy)
import Data.Maybe (fromMaybe)
import Data.Ord (comparing)

import Data.Array (Array, array, (!))

import Data.Map (Map)
import qualified Data.Map as M

import Data.Set (Set)
import qualified Data.Set as S

type Time = Int

type LogProbability = Double

data HiddenMarkovModel state observation = HMM
  { universe :: Set state
  , initial :: Map state LogProbability
  , transition :: Map state (Map state LogProbability)
  , emission :: Map observation (Map state LogProbability)
  , observation :: Map Time observation
  , timeRange :: (Time, Time)
  }
  deriving (Read, Show, Eq, Ord)

viterbi
  :: forall state observation . (Ord state, Ord observation)
  => HiddenMarkovModel state observation
  -> Maybe (LogProbability, Array Time state)
viterbi hmm
  | logProbability > log 0 = Just (logProbability, path)
  | otherwise = Nothing
  where

    t0, t1 :: Time
    (t0, t1) = timeRange hmm

    init :: state -> LogProbability
    init s = fromMaybe (log 0) $ do
                M.lookup s (initial hmm)

    emit :: Time -> state -> LogProbability
    emit t s = case M.lookup t (observation hmm) of
                Just o -> case M.lookup o (emission hmm) of
                  Just e -> fromMaybe (log 0) (M.lookup s e)
                  Nothing -> log 0
                Nothing -> init s

    prev :: Time -> state -> LogProbability
    prev t s = fst (trellis ! (t - 1) M.! s)

    trans :: state -> state -> LogProbability
    trans a b = fromMaybe (log 0) $ do
                m <- M.lookup a (transition hmm)
                M.lookup b m

    trellis :: Array Time (Map state (LogProbability, state))
    trellis = array (timeRange hmm) $
        [ (t, if t == t0

              then
                let p a = emit t a + init a
                in  M.fromList
                      [ (a, (p a, error "Viterbi.viterbi: unused backlink was used"))
                      | a <- S.toList (universe hmm)
                      ]

              else
                let p a b = emit t a + trans a b + prev t b
                in  M.fromList
                      [ (a, best
                          [ (p a b, b)
                          | b <- M.keys (transition hmm M.! a)
                          ])
                      | a <- M.keys (transition hmm)
                      ])

        | t <- [t0 .. t1]
        ]

    logProbability :: LogProbability
    finalState :: state
    (logProbability, finalState) = best
        [ fromMaybe (log 0, a)
        . fmap (second (const a))
        . M.lookup a
        . (! t1)
        $ trellis
        | a <- S.toList (universe hmm)
        ]

    path :: Array Time state
    path = array (timeRange hmm) $
        (t1, finalState) :
        [ (t, snd $ trellis ! (t + 1) M.! (path ! (t + 1)))
        | t <- [t0 .. t1 - 1]
        ]

    best :: [(LogProbability, state)] -> (LogProbability, state)
    best = maximumBy (comparing fst)

    second :: (b -> c) -> (a, b) -> (a, c)
    second f (a, b) = (a, f b)

viterbis
  :: forall state observation . (Ord state, Ord observation)
  => HiddenMarkovModel state observation
  -> [(LogProbability, [state])]
viterbis hmm = map (second reverse) (prune paths)
  where

    t0, t1 :: Time
    (t0, t1) = timeRange hmm

    init :: state -> LogProbability
    init s = fromMaybe (log 0) $ do
                M.lookup s (initial hmm)

    emit :: Time -> state -> LogProbability
    emit t s = case M.lookup t (observation hmm) of
                Just o -> case M.lookup o (emission hmm) of
                  Just e -> fromMaybe (log 0) (M.lookup s e)
                  Nothing -> log 0
                Nothing -> init s

    prev :: Time -> state -> [(LogProbability, [state])]
    prev t s = trellis ! (t - 1) M.! s

    trans :: state -> state -> LogProbability
    trans a b = fromMaybe (log 0) $ do
                m <- M.lookup a (transition hmm)
                M.lookup b m

    trellis :: Array Time (Map state [(LogProbability, [state])])
    trellis = array (timeRange hmm) $
        [ (t, if t == t0

                then
                    M.fromList
                      [ (a, prune [(emit t a + init a, [])])
                      | a <- S.toList (universe hmm)
                      ]

                else
                    M.fromList
                      [ (a, merges
                          [ prune
                              [ (p + q, b : path)
                              | (q, path) <- prev t b
                              ]
                          | b <- M.keys (transition hmm M.! a)
                          , let p = emit t a + trans a b
                          ])
                      | a <- S.toList (universe hmm)
                      ])

        | t <- [t0 .. t1]
        ]

    paths :: [(LogProbability, [state])]
    paths
        = merges
        [ map (second (a :))
        . fromMaybe []
        . M.lookup a
        . (! t1)
        $ trellis
        | a <- S.toList (universe hmm)
        ]

    prune :: [(LogProbability, [state])] -> [(LogProbability, [state])]
    prune = takeWhile ((log 0 <) . fst)

    second :: (b -> c) -> (a, b) -> (a, c)
    second f (a, b) = (a, f b)

merges :: Ord a => [[(a, b)]] -> [(a, b)]
merges = foldr merge []

merge :: Ord a => [(a, b)] -> [(a, b)] -> [(a, b)]
merge xs [] = xs
merge [] ys = ys
merge xs0@(x@(xo, _):xs) ys0@(y@(yo, _):ys)
  | xo >= yo  = x : merge xs ys0
  | otherwise = y : merge xs0 ys

# 4.2 Fever

import qualified Data.Map as M
import qualified Data.Set as S

import Viterbi

data State = Healthy | Fever
  deriving (Read, Show, Eq, Ord, Enum, Bounded)

data Observation = Normal | Cold | Dizzy
  deriving (Read, Show, Eq, Ord, Enum, Bounded)

hmm :: HiddenMarkovModel State Observation
hmm = HMM
  { universe = S.fromList [Healthy, Fever]
  , initial = M.fromList [(Healthy, log 0.6), (Fever, log 0.4)]
  , transition = M.fromList
      [ (Healthy, M.fromList [ (Healthy, log 0.7), (Fever, log 0.4) ])
      , (Fever,   M.fromList [ (Healthy, log 0.3), (Fever, log 0.6) ])
      ]
  , emission = M.fromList
      [ (Normal,  M.fromList [ (Healthy, log 0.5), (Fever, log 0.1) ])
      , (Cold,    M.fromList [ (Healthy, log 0.4), (Fever, log 0.3) ])
      , (Dizzy,   M.fromList [ (Healthy, log 0.1), (Fever, log 0.6) ])
      ]
  , observation = M.fromList
      [ (1, Normal)
      , (2, Cold)
      , (3, Dizzy)
      ]
  , timeRange = (1, 3)
  }

main :: IO ()
main = do
  print (viterbi hmm)
  mapM_ print (take 10 $ viterbis hmm)

# 4.3 Trigram

import Prelude hiding (init)

import Data.Array (elems)

import qualified Data.Map as M
import Data.Map (Map)

import Viterbi

data State a = Trigram a a a
  deriving (Read, Show, Eq, Ord)

type Observation a = a

mkhmm
  :: Ord a
  => Map (State a) LogProbability
  -> Map (State a) (Map (State a) LogProbability)
  -> Map Time (Observation a)
  -> HiddenMarkovModel (State a) (Observation a)
mkhmm init trans obs = HMM
  { universe = M.keysSet trans
  , initial = init
  , transition = trans
  , emission = fmap normalize $ M.fromListWith M.union
      [ (c, M.fromList [ (a, 1) ])
      | a@(Trigram _ c _) <- M.keys trans
      ]
  , observation = obs
  , timeRange = (minimum (M.keysSet obs), maximum (M.keysSet obs))
  }

normalize
  :: (Ord state)
  => Map state Int
  -> Map state LogProbability
normalize m = fmap (const p) m
  where
    p = log (1 / fromIntegral (sum (M.elems m)))

trigrams :: [a] -> [State a]
trigrams (a:ds@(b:c:_)) = Trigram a b c : trigrams ds
trigrams _ = []

chain :: Ord a => [a] -> Map a (Map a LogProbability)
chain ts = t
  where
    ps = zip ts (drop 1 ts)
    go (a, b) m = case M.lookup a m of
      Nothing -> M.insert a (M.singleton b 1) m
      Just ma -> M.insert a (M.insertWith (+) b 1 ma) m
    ti = foldr (.) id (map go ps) M.empty
    t = M.map (\m ->
        let s = fromInteger (sum $ M.elems m)
        in  M.map (\i -> log $ fromInteger i / s) m) ti

transpose
  :: (Ord a, Ord b)
  => Map a (Map b c)
  -> Map b (Map a c)
transpose m
  = M.fromListWith M.union
      [ (b, M.fromList [(a, c)])
      | (a, bc) <- M.toList m
      , (b, c) <- M.toList bc
      ]

main :: IO ()
main = do
  corpus <- readFile "corpus.txt"
  template <- readFile "template.txt"
  let init = normalize $ M.unionsWith (+)
          [ M.singleton a 1
          | a <- trigrams corpus
          ]
      trans = transpose . chain . trigrams $ corpus
      obs = M.fromList
          . filter ((' ' /=) . snd)
          . zip [1..]
          . filter ('\n' /=)
          $ template
      hmm = mkhmm init trans obs
      printViterbi (p, ts) =
          putStrLn $ show p ++ "\t" ++ [ c | Trigram _ c _ <- ts ]
  print (M.size trans)
  case viterbi hmm of
    Nothing -> return ()
    Just (p, ts) -> printViterbi (p, elems ts)
  mapM_ printViterbi (take 10 $ viterbis hmm)