E-step of ML estimation of hidden Markov models

updateE.jl

function updateE!(hmm::HMM,
                  Y::AbstractMatrix,     # shape: (D, T)
                  α::Matrix{Float64},   # shape: (K, T)
                  β::Matrix{Float64},   # shape: (K, T)
                  γ::Matrix{Float64},   # shape: (K, T)
                  ξ::Array{Float64, 3}, # shape: (K, K, T-1)
                  B::Matrix{Float64})    # shape: (K, T)
    const D, T = size(Y)

    # Gaussian prob.
    for k=1:hmm.K
        gauss = MvNormal(hmm.μ[:,k], hmm.Σ)
        for t=1:T
            B[k,t] = pdf(gauss, Y[:,t])
        end
    end
    
    c = Array(Float64, T)

    # forward recursion
    α[:,1] = hmm.π .* B[:,1]
    c[1] = sum(α[:,1])
    α[:,1] /= c[1]
    for t=2:T
        α[:,t] = (hmm.A' * α[:,t-1]) .* B[:,t]
        c[t] = sum(α[:,t])
        α[:,t] /= c[t]
    end
    likelihood = sum(log(c))
    @assert !any(isnan(α))

    # backword recursion
    β[:,T] = 1.0
    for t=T-1:-1:1
        for i=1:hmm.K
            s = 0.0
            for j=1:hmm.K
                s += hmm.A[i,j] * B[j,t+1] * β[j,t+1]
            end
            β[i,t] = s / c[t+1]
        end
    end
    @assert !any(isnan(β))

    γ = α .* β

    for t=1:T-1
        for i=1:hmm.K
            for j=1:hmm.K
                ξ[i,j,t] = α[i,t] * hmm.A[i,j] * B[j,t+1] * β[j,t+1] / c[t+1]
            end
        end
    end

    return γ, ξ, likelihood
end