Parameter inference for discrete hidden Markov model using autodial

hiddenmarkov.jl

using Random, LinearAlgebra, Distributions, StatsBase
using StaticArrays

# use representation x*exp(l) to compute the likelihood without underflow
struct Exp{T} <: Number 
    x::T
    l::T
end

Exp(x::T) where {T<:Real} = Exp(x, zero(x))
Exp{S}(x::Exp) where {S} = Exp(S(x.x), S(x.l))

Base.show(io::IO, x::Exp) = Base.print(io, x.x*exp(x.l - log(10)*round(Int, x.l/log(10))), "e", round(Int, x.l/log(10)))

Base.:*(x::Exp, y::Exp) = Exp(x.x*y.x, x.l + y.l)
Base.:+(x::Exp, y::Exp) = x.l==y.l ? Exp(x.x + y.x, x.l) : (x.x==0 ? y : Exp(x.x*y.x*exp(max(x.l,y.l)), min(x.l,y.l)))

Base.log(x::Exp) = log(x.x) + x.l
Base.conj(x::Exp{<:Real}) = x
Base.promote_rule(::Type{<:Exp{T}}, ::Type{S}) where {S<:Real,T} = Exp{promote_rule(S, T)}
Base.promote_rule(::Type{<:Exp{T}}, ::Type{<:Exp{S}}) where {S,T} = Exp{promote_rule(S, T)}

Base.zero(x::Exp{T}) where {T} = Exp(zero(T)) 
Base.float(x::Exp) = x.x*exp(x.l)


"""
      samplefromchain(p, P, n)

Samples an n-step Markov chain with starting distribution p
with transition matrix P.
"""
function samplefromchain(p, P, n) # n is number of trans
      s = samplefrom(p)
      statevector = [s]
      for i in 1:n
            s = samplefrom(P[s, :])
            push!(statevector, s)
      end
      return statevector
end
samplefrom(p) = sample(eachindex(p), weights(p))


Random.seed!(1)
# Model
Pgenau(λ) =  λ*@SMatrix([0 1 0; 0 0 1; 1 0 0]) + (1 - λ)*(@SMatrix(ones(3,3)))/3 
H(x) = (x==3) + 1
dirac(y) = @SVector [y==1, y==2]


# Generate data
λtrue = 0.8
P = Pgenau(λtrue)
p0 = @SVector [1,0,0]
pT = (@SVector ones(3))/3
n = 300
X = (samplefromchain(p0, P, n))
Y = H.(X)


# Compute likelihood
function l(Y, λ)
    w_ = 1.7
    w = Exp(w_, -log(w_)) # == 1
    p = Exp.(@SVector(ones(3))/3)
    π = p'
    P = Exp.(Pgenau(λ))
    O = Exp.(@SMatrix [1.0 0; 1.0 0; 0 1.0])
    n = length(Y)
    for i in n:-1:2
        p = P*p
        p = (O*Exp.(1.0*dirac(Y[i]))) .* p
        p = p .* w
    end
    π*p   
end
logl(Y) = λ -> log(l(Y, λ))

# Rough maximum likelihood estimate
r = 0.7:0.01:0.9
ls = [λ => log(l(Y, λ)) for λ in r]
println("Likelihoods for λ = ", r)
println.(ls);
ll, i = findmax((last.(ls)))
println("maximum likelihood estimate: ", r[i])


# Make figure of likelihood
using Makie
sep(x) = first.(x), float.(last.(x))
fig = Figure(resolution=(1600,800))
scatter!(Axis(fig[1,1], title="data"), 1:length(Y), Y, markersize=2.0)

lines!(Axis(fig[2,1], title="log likelihood λ"), sep(ls)...)
fig 

# Make likelihood differentiable
using ForwardDiff
using ForwardDiff: Dual, partials
@inline ForwardDiff.extract_derivative(::Type{T}, y::Exp{<:Dual}) where {T} = partials(T, y.x*exp(y.l), 1)
dls = [ λ => ForwardDiff.derivative(logl(Y), λ) for λ in r]


# Add a figure for derivative
ax2 = Axis(fig[3,1], title="∇ log likelihood λ")
lines!(ax2, sep(dls)...)
lines!(ax2, r, 0 .* r, linestyle=:dot)
fig


# Use my favourite sampler to take prior and uncertainty into account
using ZigZagBoomerang
using SparseArrays
# minus derivative of log likelihood + log prior
∇ϕ(x, _, Y) = -partials.(log(l(Y, Dual(x[], 1.0))), 1) #uniform prior 

t0 = 0.0
x0 = [0.75]
θ0 = [1.0]
c = [100.0]
Γ = sparse(1.0ones(1,1))
Z = ZigZag(Γ, x0)
T = 100.0
println("sampling")
trace, _, acc = @time spdmp(∇ϕ, t0, x0, θ0, T, c, Z, Y; adapt=true)
dt = 0.5
ts, xs = sep(collect(discretize(trace, dt))) # use for stats
ts_, xs_ = sep(collect(trace)) # nicer for plotting

println("Bayes estimate: ", mean(xs)[], " ± ", std(xs)[])

# Update figure with posterior trace
ax3 = Axis(fig[4,1], title="trace and truth")
lines!(ax3, ts_, first.(xs_), color=:black)
lines!(ax3, [0.0, T], [λtrue, λtrue], linewidth=2.0, color=:orange)
fig