Viterbi algorithm example

viterbi.js

const obs = ["normal", "cold", "dizzy"]
const states = ["Healthy", "Fever"]
const start_p = {
    Healthy: 0.6,
    Fever: 0.4,
}

const trans_p = {
    Healthy: {
        Healthy: 0.7,
        Fever: 0.3,
    },
    Fever: {
        Healthy: 0.4,
        Fever: 0.6,
    },
}

const emit_p = {
    Healthy: {
        normal: 0.5,
        cold: 0.4,
        dizzy: 0.1,
    },
    Fever: {
        normal: 0.1,
        cold: 0.3,
        dizzy: 0.6,
    },
}

function viterbi(obs, states, start_p, trans_p, emit_p) {
    let last_states = undefined
    return obs.map((ob) => {
        if (last_states == undefined) {
            last_states = {}
            for (const state of states) {
                last_states[state] = start_p[state] * emit_p[state][ob]
            }
        } else {
            let current_states = {}

            for (const state of states) {
                current_states[state] = Math.max(
                  ...states.map(
                    (new_state) =>
                      last_states[new_state] *
                      trans_p[state][new_state] *
                      emit_p[new_state][ob]
                  )
                )
            }

            last_states = current_states
        }

        let max = 0
        let maxState = undefined

        for (const state in last_states) {
            const prob = last_states[state]
            if (prob > max) {
                max = prob
                maxState = state
            }
        }

        return maxState
    })
}

const path = viterbi(obs, states, start_p, trans_p, emit_p)
console.log(path)