Viterbi in Scala

Viterbi.scala

package dk.linnk.hmm

object Viterbi {
  type State = String
  type Observation = String
  type Probability = Double
  type ProbabilityMap = (State, State) => Probability
  type EmissionMap = (State, Observation) => Probability
  type ProbabilityPath = (Probability, List[State])

  def viterbi(observations: List[Observation],
              states: List[State],
              start: State => Probability,
              transition: ProbabilityMap,
              emissions: EmissionMap): ProbabilityPath = {

    def probability(p: ProbabilityPath) = p._1

    def mostLikelyPathFrom(state: State, time: Int): ProbabilityPath = {
      val emission = emissions(state, observations(time))
      time match {
        case 0 =>
          // (probability that were in the initial state) times
          // (probability of observing the initial observation from the initial state)
          (start(state) * emission, List(state))
        case _ =>
          val (prob, path) = states map { (state) =>
            val (prob, path) = mostLikelyPathFrom(state, time - 1)
            val prevState = path.head
            // (probability of the previous state) times
            // (probability of moving from previous state to this state)
            (prob * transition(prevState, state), path)
          } maxBy probability
          // (probability of observing the current observation from this state) times
          // (probability of the maximizing state)
          (emission * prob, state :: path)
      }
    }

    val (prob, path) = states map { (state) =>
      mostLikelyPathFrom(state, observations.size - 1)
    } maxBy probability

    (prob, path.reverse)
  }
}

object VibertiTest extends App {
  val states = List("Healthy", "Fever")
  val observations = List("normal", "cold", "dizzy")
 
  val start_probability : String => Double = {
    case "Healthy" => 0.6
    case "Fever"   => 0.4
  }
 
  val transition_probability : (String, String) => Double = {
    case ("Healthy", "Healthy") => 0.7
    case ("Healthy", "Fever")   => 0.3
    case ("Fever", "Healthy")   => 0.4
    case ("Fever", "Fever")     => 0.6
  }
 
  val emission_probability : (String, String) => Double = {
    case ("Healthy", "normal") => 0.5
    case ("Healthy", "cold")   => 0.4
    case ("Healthy", "dizzy")  => 0.1
    case ("Fever", "normal")   => 0.1
    case ("Fever", "cold")     => 0.3
    case ("Fever", "dizzy")    => 0.6
  }

  println(Viterbi.viterbi(observations,
                          states,
                          start_probability,
                          transition_probability,
                          emission_probability))
}