Bootstrap filtering and smoothing

This tutorial describes how to perform filtering and smoothing in a the probabilistic state-space model given by

\[\begin{aligned} x_0 &\sim \mathcal{N}(\mu_0 ,\Sigma_0), \\ x_n \mid x_{n-1} &\sim \mathcal{N}(\Phi x_{n-1}, Q),\\ z_n &= C x_n, \end{aligned}\]

subject to the measurements given by

\[y_n \mid x_n \sim \mathcal{N}(0, e^{z_n}).\]

Setting up the environment and generating some data

using MarkovKernels
using LinearAlgebra, Plots, Random
using IterTools
import StatsBase: wsample

rng = MersenneTwister(1991)

function sample(rng, init, K, nstep)
    it = Iterators.take(iterated(z -> rand(rng, K, z), rand(rng, init)), nstep + 1)
    return mapreduce(permutedims, vcat, collect(it))
end

# time grid
m = 50
T = 5
ts = collect(LinRange(0, T, m))
dt = T / (m - 1)

# define transtion kernel
λ = 1.0
Φ = exp(-λ * dt) .* [1.0 0.0; -2*λ*dt 1.0]
Q = I - exp(-2 * λ * dt) .* [1.0 -2*λ*dt; -2*λ*dt 1+(2*λ*dt)^2]
fw_kernel = NormalKernel(Φ, Q)

# initial distribution
init = Normal(zeros(2), 1.0I(2))

# sample state
xs = sample(rng, init, fw_kernel, m - 1)

# output kernel
σ = 5.0
C = σ / sqrt(2) * [1.0 -1.0]

output_kernel = DiracKernel(C)

variance(x) = fill(exp.(x)[1], 1, 1)
m_kernel = compose(NormalKernel(zeros(1, 1), variance), output_kernel)

# sample output
zs = mapreduce(z -> rand(rng, output_kernel, xs[z, :]), vcat, 1:m)
ys = mapreduce(z -> rand(rng, m_kernel, xs[z, :]), vcat, 1:m)

measurement_plt = scatter(ts, ys, label = "measurements", color = "black")
Example block output

Implementing a bootstrap filter

function bootstrap_filter(
    rng::AbstractRNG,
    ys::AbstractVecOrMat,
    init::AbstractDistribution,
    fw_kernel::AbstractMarkovKernel,
    m_kernel::AbstractMarkovKernel,
    K::Integer,
)

    # initialize
    X = [rand(rng, init) for k in 1:K]
    P = ParticleSystem(zeros(K), X)
    loglike = 0.0
    L = Likelihood(m_kernel, ys[1, :])
    loglike_incr = posterior_and_loglike!(P, L)
    loglike = loglike + loglike_incr
    resample!(rng, P)

    Ps = [P]
    sizehint!(Ps, size(ys, 1))

    for m in 2:size(ys, 1)
        L = Likelihood(m_kernel, ys[m, :])
        P = predict(rng, P, fw_kernel)
        loglike_incr = posterior_and_loglike!(P, L)
        loglike = loglike + loglike_incr
        resample!(rng, P)
        push!(Ps, P)
    end

    return Ps, loglike
end

function resample!(rng::AbstractRNG, P::ParticleSystem{T,U,<:AbstractVector}) where {T,U}
    idx = wsample(rng, eachindex(logweights(P)), weights(P), nparticles(P))
    logweights(P)[:] .= zero(logweights(P))
    particles(P)[:] .= particles(P)[idx]
end

function predict(
    rng::AbstractRNG,
    P::ParticleSystem{T,U,<:AbstractVector},
    K::AbstractMarkovKernel,
) where {T,U}
    X = [rand(rng, K, particles(P)[i]) for i in eachindex(particles(P))]

    return ParticleSystem(copy(logweights(P)), X)
end

Computing the filtered state estimates

K = 500
Pfilt, loglike_filt = bootstrap_filter(rng, ys, init, fw_kernel, m_kernel, K)

Xfilt = mapreduce(permutedims, vcat, particles.(Pfilt))

state_filt_plt = plot(
    ts,
    xs,
    layout = (2, 1),
    xlabel = "t",
    labels = ["x1" "x2"],
    title = ["Latent Gauss-Markov process" ""],
)
for k in 1:K
    scatter!(
        ts,
        mapreduce(permutedims, vcat, Xfilt[:, k]),
        markersize = 1,
        color = "red",
        alpha = 0.025,
        label = "",
    )
end
state_filt_plt
Example block output

Computing the filtered output estimates

bf_output_filt = [marginalize(Pfilt[i], output_kernel) for i in eachindex(Pfilt)]
Zfilt = getindex.(mapreduce(permutedims, vcat, particles.(bf_output_filt)), 1)

output_filt_plt = plot(ts, zs, label = "output", xlabel = "t", title = "log-variance")
scatter!(ts, Zfilt, markersize = 1, color = "red", alpha = 0.01, label = "")
output_filt_plt
Example block output

Implementing a bootstrap smoother

function bootstrap_smoother(
    rng::AbstractRNG,
    ys::AbstractVecOrMat,
    init::AbstractDistribution,
    fw_kernel::AbstractMarkovKernel,
    m_kernel::AbstractMarkovKernel,
    K::Integer,
)

    # initialize
    X = permutedims([rand(rng, init) for k in 1:K])
    P = ParticleSystem(zeros(K), X)
    loglike = 0.0
    L = Likelihood(m_kernel, ys[1, :])
    loglike_incr = posterior_and_loglike!(P, L)
    loglike = loglike + loglike_incr
    resample!(rng, P)

    for m in 2:size(ys, 1)
        L = Likelihood(m_kernel, ys[m, :])

        P = predict(rng, P, fw_kernel)
        loglike_incr = posterior_and_loglike!(P, L)

        loglike = loglike + loglike_incr
        resample!(rng, P)
    end

    return P, loglike
end

function resample!(rng::AbstractRNG, P::ParticleSystem{T,U,<:AbstractMatrix}) where {T,U}
    idx = wsample(rng, eachindex(logweights(P)), weights(P), nparticles(P))
    logweights(P)[:] .= zero(logweights(P))
    particles(P)[:, :] .= particles(P)[:, idx]
end

function predict(
    rng::AbstractRNG,
    P::ParticleSystem{T,U,<:AbstractMatrix},
    K::AbstractMarkovKernel,
) where {T,U}

    X = [rand(rng, K, particles(P)[end, i]) for i in 1:nparticles(P)]

    return ParticleSystem(logweights(P), vcat(particles(P), permutedims(X)))
end
nothing # hide

Computing the smoothed state estimates

Psmooth, loglike_smooth = bootstrap_smoother(rng, ys, init, fw_kernel, m_kernel, K)

Xsmooth = particles(Psmooth)

bf_output_smooth = marginalize(Psmooth, output_kernel)
Ysmooth = getindex.(particles(bf_output_smooth), 1)

state_smooth_plt = plot(
    ts,
    xs,
    layout = (2, 1),
    xlabel = "t",
    labels = ["x1" "x2"],
    title = ["Latent Gauss-Markov process" ""],
)
for k in 1:K
    plot!(
        ts,
        mapreduce(permutedims, vcat, Xsmooth[:, k]),
        color = "green",
        alpha = 0.025,
        label = "",
    )
end
state_smooth_plt
Example block output

Computing the smoothed output estimate

bf_output_smooth = marginalize(Psmooth, output_kernel)
Zsmooth = getindex.(particles(bf_output_smooth), 1)

output_smooth_plt = plot(ts, zs, label = "output", xlabel = "t", title = "log-variance")
plot!(ts, Zsmooth, color = "green", alpha = 0.025, label = "")
output_smooth_plt
Example block output