Viterbi.hs

-- | Viterbi algorithm in Haskell, legacy of my course on bio-informatics.
--
module Viterbi where

import Data.Ord (comparing)
import Data.Maybe (fromJust)
import Data.List (sort, maximumBy)
import Data.Map (Map)
import qualified Data.Map as M

viterbi :: (Eq a, Ord a, Eq b, Ord b)
        => [a]                -- ^ States.
        -> (a -> a -> Float)  -- ^ State transitions.
        -> (a -> b -> Float)  -- ^ Emission probability.
        -> [b]                -- ^ Observed emissions.
        -> (Float, [a])       -- ^ Most likely sequence.
viterbi states transitionProbablity emissionProbablity emissions =
    let initial = M.fromList $ flip map states $ \s -> (s, (1, []))
        (s, tr) = viterbi' initial emissions
    in (s, reverse tr)
  where
    viterbi' prev []     = maximumBy (comparing fst) $ map snd $ M.toList prev
    viterbi' prev (x:xs) =
        let probablities = M.fromList $ zip states $ map probablity states
            probablity s = let possibilies = map (\x -> comeFrom x s) states
                               (p, tr) = maximumBy (comparing fst) possibilies
                           in (p * emissionProbablity s x, s : tr)
            comeFrom p s = let (prevP, prevTR) = fromJust (M.lookup p prev)
                           in (prevP * transitionProbablity p s, prevTR)
        in viterbi' probablities xs

-- What follows below is a concrete example.

data State = Zero | Background | Promoter
           deriving (Eq, Ord, Show)

states = [Zero, Background, Promoter]

tProb Zero Background = 0.5
tProb Zero Promoter = 0.5
tProb Background Background = 0.85
tProb Background Promoter = 0.15
tProb Promoter Promoter = 0.75
tProb Promoter Background = 0.25
tProb _ _ = 0.0

data Emission = A | C | G| T
              deriving (Eq, Ord, Show)

emissions = [A, C, G, T]

eProb Zero _ = 0.0
eProb Background _ = 0.25
eProb Promoter A = 0.15
eProb Promoter C = 0.42
eProb Promoter G = 0.30
eProb Promoter T = 0.13