Updating model parameters in a Gibbs sampler with CSMC

[bonStats]

Hi there,

I am helping to implement a Gibbs sampler for a state space model where

• the model parameters are updated with a MH step,
• the latent states are updated with PG or PGAS,

with AdvancedPS.jl, thanks!

We have been running into problems with convergence, so I started to play with a simpler model/implementation to try and find the problem.

I adapted the Gaussian SSM example to manually sample the latent states and have found the following problem.

The below code implements a sampler where the model parameters increase every iteration (of course this won't target a distribution we care about) but does lead to finding that the model carried by AdvancedPS.PGState isn't updating. Does this mean the AbstractMCMC.step function ignores the model argument and uses the model from the state?

using AdvancedPS
using Random
using Distributions
using Plots
using AbstractMCMC
using Random123
using SSMProblems

"""
    plot_update_rate(update_rate, N)

Plot empirical update rate against theoretical value
"""
function plot_update_rate(update_rate::AbstractVector{Float64}, Nₚ::Int)
    plt = plot(
        update_rate;
        label=false,
        ylim=[0, 1],
        legend=:bottomleft,
        xlabel="Iteration",
        ylabel="Update rate",
    )
    return hline!(plt, [1 - 1 / Nₚ]; label="N: $(Nₚ)")
end

"""
    update_rate(trajectories, N)

Compute latent state update rate
"""
function update_rate(particles::AbstractMatrix{Float64}, Nₛ)
    return sum(abs.(diff(particles; dims=2)) .> 0; dims=2) / Nₛ
end

struct GaussianPrior{T<:Real} <: SSMProblems.StatePrior
    σ::T
end

function SSMProblems.distribution(proc::GaussianPrior)
    return Normal(0, proc.σ)
end

struct LinearGaussianDynamics{AT<:Real,QT<:Real} <: SSMProblems.LatentDynamics
    a::AT
    q::QT
end

function SSMProblems.distribution(dyn::LinearGaussianDynamics, ::Int, state)
    return Normal(dyn.a * state, dyn.q)
end

struct StochasticVolatility <: SSMProblems.ObservationProcess end

function SSMProblems.distribution(::StochasticVolatility, ::Int, state)
    return Normal(0, exp(state / 2))
end

function LinearGaussianStochasticVolatilityModel(a, q)
    prior = GaussianPrior(q)
    dyn = LinearGaussianDynamics(a, q)
    obs = StochasticVolatility()
    return SSMProblems.StateSpaceModel(prior, dyn, obs)
end

rng = Random.MersenneTwister(1234)
true_model = LinearGaussianStochasticVolatilityModel(0.9, 0.5)
_, x, y = sample(rng, true_model, 200);

plot(x; label="x", xlabel="t")

plot(y; label="y", xlabel="t")

## Model parameters are fixed: use sample
model = AdvancedPS.TracedSSM(true_model, y)
pg = AdvancedPS.PGAS(20, AdvancedPS.ResampleWithESSThreshold(0.5))
#pg = AdvancedPS.PG(20, AdvancedPS.ResampleWithESSThreshold(0.5))
chains = sample(rng, model, pg, 200; progress=false);

##

## Model parameters update each iteration: DIY with step
_, state0 = AbstractMCMC.step(rng, model, pg)
vals = Vector{typeof(state0)}(undef, 200)
last_val = state0

model_inc = AdvancedPS.TracedSSM(LinearGaussianStochasticVolatilityModel(0.9, 0.5), y) # to test

for ss in 1:200
    model_inc = AdvancedPS.TracedSSM(LinearGaussianStochasticVolatilityModel(0.9+ss/500, 0.5+ss/500), y) # to test
    _, state_ss = AbstractMCMC.step(rng, model_inc, pg, last_val)
    vals[ss] = last_val = state_ss
end

# PROBLEM...

# every state carries the current model parameter values
vals[10].trajectory.model.model.dyn.a == vals[20].trajectory.model.model.dyn.a
# true
# they shouldn't be equal, I've explicity changed them with the model_inc
# is the model carries by last_val overriding the explicit model passed as the second argument?

But of course, I could just be doing something not recommend or plain wrong!