viterbi.rb

viterbi.rb

# This script is a simple implementation of the famous Viterbi algorithm.
# Viterbi can find the best path of states through an HMM (Hidden Markov Model),
# given a sequence of observations. In the future, the HMM class could be
# extended to also support other algorithms (e.g. Forward or Baum-Welch).
#
# https://de.wikipedia.org/wiki/Hidden_Markov_Model
# https://en.wikipedia.org/wiki/Viterbi_algorithm
#
class HMM

  # Initialize with given HMM model. Each state consists of a name, emission
  # probabilities, transition probabilities and an initial probability.
  # Additional basic checks to ensure probability sums of 1 are in place.
  def initialize(states)

    # Guard that emissions always sum up to 1
    if states.any? { |s| s[:emissions].values.sum != 1 }
      raise "Probabilities of emissions do no sum up to 1."
    end

    # Guard that transitions always sum up to 1
    if states.any? { |s| s[:transitions].values.sum != 1 }
      raise "Probabilities of transitions do no sum up to 1."
    end

    # Keep HMM model as instance variable
    @states = states
  end

  # Finds the most likely path through the provided model, given a sequence
  # of observations. The algorithm works by first calculating the best states
  # for each observations, building upon the previous probabilities and then
  # use backtracking to find the best path through the states.
  def viterbi_path(observations)
    values = []

    # Initialize values array for each state with the first observation
    # E.g. [{ H: {prob: 0.15, prev: nil}, L: {prob: 0.1, prev: nil}}]
    values << @states.to_h do |state|
      [state[:name], {
        prob: state[:initial] * state[:emissions][observations.first],
        prev: nil
      }]
    end

    # Run the algorithm for all remaining observations
    observations.drop(1).each.with_index(1) do |observation, index|

      # For each observation, iterate through all states
      values << @states.to_h do |state|

        # Compute the best previous state
        prev_state = @states.max_by do |prev_state|
          values[index - 1][prev_state[:name]][:prob] *
          prev_state[:transitions][state[:name]]
        end

        # Compute the probability of the best previous state
        max_prob = (
          values[index - 1][prev_state[:name]][:prob] *
          prev_state[:transitions][state[:name]]
        )

        # Set new probability and backward pointer
        [state[:name], {
          prob: max_prob * state[:emissions][observation],
          prev: prev_state[:name]
        }]
      end
    end

    # For tracking back the best path, we have to find the best last state
    last_name, last_values = values.last.max_by { |_, x| x[:prob] }

    # Save the best path through the states.
    # Initialized with the name of the best last state.
    best_path = [last_name]

    # The pointer to the best previous state
    previous = last_values[:prev]

    # Iterate backward, following the pointer
    values.drop(1).reverse.each do |value|
      best_path.prepend(value[previous][:prev])
      previous = value[previous][:prev]
    end

    # Return optimal path and probability of that path
    [best_path, last_values[:prob]]
  end
end

Use it like that:

# Define a HMM model
hmm = HMM.new([
  {
    name: :Healthy,
    emissions: { normal: 0.5, cold: 0.4, dizzy: 0.1 },
    transitions: { Healthy: 0.7, Fever: 0.3 },
    initial: 0.6
  }, {
    name: :Fever,
    emissions: { normal: 0.1, cold: 0.3, dizzy: 0.6 },
    transitions: { Healthy: 0.4, Fever: 0.6 },
    initial: 0.4
  }
])

# Compute best path and its probability, given the observations
path, prob = hmm.viterbi_path([:normal, :cold, :dizzy])
puts "Likeliest path: #{path.join(', ')} with probability: #{prob}"