Viterbi decoder in Julia

viterbi.jl

typealias State Uint
typealias Observation Uint


function decode(x::Vector{Observation}, transition_matrix::Matrix{Float64},
                observation_matrix::Matrix{Float64})
    n_observations::Uint = uint(length(x))
    n_states::Uint = uint(size(transition_matrix)[1])
    @assert(size(transition_matrix) == size(transpose(transition_matrix)))

    log_p_transition(s_next::State, s::State) = log(transition_matrix[s, s_next])
    log_p_observation(obs::Observation, s::State) = log(observation_matrix[s, obs])

    log_p_y = zeros(n_states, n_observations)
    prev_state = zeros(State, n_states, n_observations)

    # Initialize the likelihoods of the first state (assuming a uniform
    # prior over states)
    for s in uint(1):n_states
        log_p_y[s, 1] = log_p_observation(x[1], s)
        prev_state[s, 1] = s
    end

    for i in uint(2):n_observations
        for s in uint(1):n_states
            # Determine the log likelihood of the sequence terminating in state
            # s.
            log_p_prev, prev_state[s, i] = findmax([(log_p_transition(s, s_prev) +
                                          log_p_y[s_prev, i-1])
                                         for s_prev in uint(1):n_states])
            log_p_y[s, i] = log_p_observation(x[i], s) + log_p_prev
        end
    end

    # Determine the maximum likelihood sequence by finding the maximum
    # likelihood end state and following the sequence back.
    ml_sequence = zeros(State, n_observations)
    ml_sequence[n_observations] = indmax(log_p_y[:, n_observations])
    for i in (n_observations-1):-1:uint(1)
        ml_sequence[i] = prev_state[ml_sequence[i+1], i+1]
    end

    ml_sequence
end