F# Implementation of Viterbi

viterbi.fsx

(* Nick Heiner - Info 2950 PS9 *)

(* Viterbi algorithm, as described here: http://people.ccmr.cornell.edu/~ginsparg/INFO295/vit.pdf 

  priorProbs: prior probability of a hidden state occuring
  transitions: probability of one hidden state transitioning into another
  emissionProbs: probability of a hidden state emitting an observed state
  observation: a sequence of observed states
  hiddens: a list of all possible hidden states

  Returns: probability of most likely path * hidden state list representing the path

*)
let viterbi (priorProbs : 'hidden -> float) (transitions : ('hidden * 'hidden) -> float) (emissionProbs : (('observed * 'hidden) -> float))
  (observation : 'observed []) (hiddens : 'hidden list) : float * 'hidden list =

  (* Referred to as v_state(time) in the notes *)
  (* Probability of the most probable path ending in state at time *)
  let rec mostLikelyPathProb (state : 'hidden) (time : int) : float * 'hidden list =
    let emission = emissionProbs (observation.[time], state)
    match time with 
      (* If we're at time 0, then just use the emission probability and the prior probability for this state *)
      | 0 -> emission * priorProbs state, [state]

      (* If we're not at time 0, then recursively look for the most likely path ending at this time *)
      | t when t > 0 ->
          let prob, path = Seq.maxBy fst (seq { for hiddenState in hiddens -> 
                                                (* Recursively look for most likely path at t - 1 *)
                                                let prob, path = mostLikelyPathProb hiddenState (time - 1)
                                                (* Rate each path by how likely it is to transition into the current state *)
                                                transitions (List.head path, state) * prob, path})
          emission * prob, state::path

      (* If time is < 0, then throw an error *)
      | _ -> failwith "time must be >= 0"

  (* Look for the most likely path that ends at t_max *)
  let prob, revPath = Seq.maxBy fst (seq { for hiddenState in hiddens -> mostLikelyPathProb hiddenState ((Array.length observation) - 1)}) 
  prob, List.rev revPath

(* observable states *)
type signal = Yes | No

(* hidden states *)
type reality = Rain | Dry
let states = [Rain; Dry]

let rec emissionProbs = function
  | (Yes, Rain) -> 0.6
  | (Yes, Dry) -> 1. - emissionProbs (Yes, Rain)
  | (No, Rain) -> 1. - emissionProbs (No, Dry)
  | (No, Dry) -> 0.8

(* Probability of a transition between (prev, next) *)
let rec transitions = function
  | (Rain, Rain) -> 0.65
  | (Rain, Dry) -> 1. - transitions (Rain, Rain)
  | (Dry, Rain) -> 0.25 
  | (Dry, Dry) -> 1. - transitions (Dry, Rain)

let priorProbs (_ : reality) = 0.5

let observed = [| Yes; Yes; No; Yes; No; No; Yes |]

viterbi priorProbs transitions emissionProbs observed states
(* Outputs (0.0011664, [Dry; Dry; Dry; Dry; Dry; Dry; Dry]), which I must admit I'm a bit skeptical about. *)

(* Testing from course notes: http://people.ccmr.cornell.edu/~ginsparg/INFO295/vit.pdf *)
type observed = R | B
type hidden = One | Two | Three
let hiddens = [One; Two; Three]
let priors _ = ((float)1/(float)3)
let emissions = function
  | (R, One) -> 0.5
  | (B, One) -> 0.5
  | (R, Two) -> ((float)1/(float)3)
  | (B, Two) -> ((float)2/(float)3)
  | (R, Three) -> ((float)3/(float)4)
  | (B, Three) -> ((float)1/(float)4)

let transitionProbs = function
  (* from, to *)
  | (One, One) -> 0.3
  | (One, Two) -> 0.6
  | (One, Three) -> 0.1
  | (Two, One) -> 0.5
  | (Two, Two) -> 0.2
  | (Two, Three) -> 0.3
  | (Three, One) -> 0.4
  | (Three, Two) -> 0.1
  | (Three, Three) -> 0.5

let observations = [| R; B; R |]

viterbi priors transitionProbs emissions observations hiddens
(* Outputs (0.01666666667, [One; Two; One]), which is correct *)

(* testing with example from wikipedia: http://en.wikipedia.org/wiki/Viterbi_algorithm#Example *)
type wikiHiddens = Healthy | Fever
let wikiHiddenList = [Healthy; Fever]
type wikiObservations = Normal | Cold | Dizzy

let wikiPriors = function
  | Healthy -> 0.6
  | Fever -> 0.4

let wikiTransitions = function
  | (Healthy, Healthy) -> 0.7
  | (Healthy, Fever) -> 0.4
  | (Fever, Healthy) -> 0.4
  | (Fever, Fever) -> 0.6

let wikiEmissions = function
  | (Cold, Healthy) -> 0.4
  | (Normal, Healthy) -> 0.5
  | (Dizzy, Healthy) -> 0.1
  | (Cold, Fever) -> 0.3
  | (Normal, Fever) -> 0.1
  | (Dizzy, Fever) -> 0.6

viterbi wikiPriors wikiTransitions wikiEmissions [| Dizzy; Normal; Cold |] wikiHiddenList
(* Outputs: (0.01344, [Fever; Healthy; Healthy]), which is correct *)