Hidden Markov Models with Gaussian observation

GaussianHMM.jl

# Hidden Markov Models

using Distributions

type HMM
    K::Int                # number of possible states
    μ::Array{Float64, 2} # shape (D, K)
    π::Array{Float64, 1} # shape (K)
    A::Array{Float64, 2}  # shape (K, K)
    Σ::Array{Float64, 3} # shape (D, D, K)
    
    function HMM(D::Int, K::Int)
        Σ = Array(Float64, D, D, K)
        for k=1:K
            Σ[:,:,k] = diagm(ones(D))
        end
        new(K,
            rand(D, K),
            ones(K) ./ K,
            ones(K, K) ./ K,
            Σ)
    end
end

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.Σ[:,:,k])
        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
    @assert !any(isnan(α))

    likelihood = sum(log(c))
    @assert !isnan(likelihood)

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

    γ = α .* β

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

    return γ, ξ, likelihood
end

function updateM!(hmm::HMM,
                  Y::AbstractMatrix,
                  γ::Matrix{Float64},   # shape: (K, T)
                  ξ::Array{Float64, 3}) # shape: (K, K, T-1)
    const D, T = size(Y)
    
    # prior
    hmm.π[:] = γ[:,1] / sum(γ[:,1])

    # observation
    # mean
    hmm.μ[:,:] = (Y*γ') ./ sum(γ, 2)'

    # covariance
    # fill!(hmm.Σ, 0.0)
    for k=1:hmm.K
        Ŷ = Y .- hmm.μ[:,k]
        hmm.Σ[:,:,k] = (γ[k,:] .* Ŷ) * Ŷ' / sum(γ[k,:])
    end

    # transition
    hmm.A[:,:] = sum(ξ, 3) ./ (sum(γ, 2))

    nothing
end

type HMMTrainingResult
    likelihoods::Vector{Float64}
    α::Matrix{Float64}
    β::Matrix{Float64}
    γ::Matrix{Float64}
    ξ::Array{Float64, 3}
end

function fit!(hmm::HMM,
              Y::AbstractMatrix;
              maxiter::Int=100,
              tol::Float64=0.0,
              verbose::Bool=false)
    const D, T = size(Y)
    
    likelihood::Vector{Float64} = zeros(1)

    α = Array(Float64, hmm.K, T)
    β = Array(Float64, hmm.K, T)
    γ = Array(Float64, hmm.K, T)
    ξ = Array(Float64, hmm.K, hmm.K, T-1)
    B = Array(Float64, hmm.K, T)

    for iter=1:maxiter
        # update expectations
        γ, ξ, score = updateE!(hmm, Y, α, β, γ, ξ, B)

        # update parameters
        updateM!(hmm, Y, γ, ξ)

        improvement = (score - likelihood[end]) / abs(likelihood[end])

        if verbose
            println("#$(iter): bound $(likelihood[end])
                    improvement: $(improvement)")
        end
        
        # check if converged
        if abs(improvement) < tol
            if verbose
                println("#$(iter) converged")
            end
            break
        end

        push!(likelihood, score)
    end
    
    # remove initial zero
    shift!(likelihood)
    
    return HMMTrainingResult(likelihood, α, β, γ, ξ)
end