Interface Changes for Use in Filtering

examples/particle-mcmc/Project.toml

@@ -0,0 +1,12 @@
+[deps]
+AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
+AdvancedMH = "5b7e9947-ddc0-4b3f-9b55-0d8042f74170"
+CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
+DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+GaussianDistributions = "43dcc890-d446-5863-8d1a-14597580bb8d"
+PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
+SSMProblems = "26aad666-b158-4e64-9d35-0e672562fa48"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
+UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"

examples/particle-mcmc/particles.jl

@@ -0,0 +1,180 @@
+using DataStructures: Stack
+
+## PARTICLES ###############################################################################
+
+abstract type AbstractParticleContainer{T} end
+
+"""
+    store!(particles, new_states, [idx])
+
+update the state component of the particle container, with optional parent indices supplied
+for use in ancestry storage.
+"""
+function store! end
+
+"""
+    reset_weights!(particles)
+
+in-place method to reset the log weights of the particle cloud to zero; typically called
+following a resampling step.
+"""
+function reset_weights! end
+
+mutable struct ParticleContainer{T,WT<:Real} <: AbstractParticleContainer{T}
+    vals::Vector{T}
+    log_weights::Vector{WT}
+end
+
+Base.collect(pc::ParticleContainer) = pc.vals
+Base.length(pc::ParticleContainer) = length(pc.vals)
+Base.keys(pc::ParticleContainer) = LinearIndices(pc.vals)
+
+# not sure if this is kosher, since it doesn't follow the convention of Base.getindex
+Base.@propagate_inbounds Base.getindex(pc::ParticleContainer, i::Int) = pc.vals[i]
+Base.@propagate_inbounds Base.getindex(pc::ParticleContainer, i::Vector{Int}) = pc.vals[i]
+
+function store!(pc::ParticleContainer, new_states, idx...; kwargs...)
+    setindex!(pc.vals, new_states, eachindex(pc))
+    return pc
+end
+
+function reset_weights!(pc::ParticleContainer{T,WT}) where {T,WT<:Real}
+    fill!(pc.log_weights, zero(WT))
+    return pc.log_weights
+end
+
+## JACOB-MURRAY PARTICLE STORAGE ###########################################################
+
+Base.append!(s::Stack, a::AbstractVector) = map(x -> push!(s, x), a)
+
+mutable struct ParticleTree{T}
+    states::Vector{T}
+    parents::Vector{Int64}
+    leaves::Vector{Int64}
+    offspring::Vector{Int64}
+    free_indices::Stack{Int64}
+
+    function ParticleTree(states::Vector{T}, M::Integer) where {T}
+        nodes = Vector{T}(undef, M)
+        initial_free_indices = Stack{Int64}()
+        append!(initial_free_indices, M:-1:(length(states) + 1))
+        @inbounds nodes[1:length(states)] = states
+        return new{T}(
+            nodes, zeros(Int64, M), 1:length(states), zeros(Int64, M), initial_free_indices
+        )
+    end
+end
+
+Base.length(tree::ParticleTree) = length(tree.states)
+Base.keys(tree::ParticleTree) = LinearIndices(tree.states)
+
+function prune!(tree::ParticleTree, offspring::Vector{Int64})
+    ## insert new offspring counts
+    setindex!(tree.offspring, offspring, tree.leaves)
+
+    ## update each branch
+    @inbounds for i in eachindex(offspring)
+        j = tree.leaves[i]
+        while (j > 0) && (tree.offspring[j] == 0)
+            push!(tree.free_indices, j)
+            j = tree.parents[j]
+            if j > 0
+                tree.offspring[j] -= 1
+            end
+        end
+    end
+end
+
+function insert!(
+    tree::ParticleTree{T}, states::Vector{T}, a::AbstractVector{Int64}
+) where {T}
+    ## parents of new generation
+    parents = getindex(tree.leaves, a)
+
+    ## ensure there are enough dead branches
+    if (length(tree.free_indices) < length(a))
+        @debug "expanding tree"
+        expand!(tree)
+    end
+
+    ## find places for new states
+    @inbounds for i in eachindex(states)
+        tree.leaves[i] = pop!(tree.free_indices)
+    end
+
+    ## insert new generation and update parent child relationships
+    setindex!(tree.states, states, tree.leaves)
+    setindex!(tree.parents, parents, tree.leaves)
+    return tree
+end
+
+function expand!(tree::ParticleTree)
+    M = length(tree)
+    resize!(tree.states, 2 * M)
+
+    # new allocations must be zero valued, this is not a perfect solution
+    tree.parents = [tree.parents; zero(tree.parents)]
+    tree.offspring = [tree.offspring; zero(tree.offspring)]
+    append!(tree.free_indices, (2 * M):-1:(M + 1))
+    return tree
+end
+
+function get_offspring(a::AbstractVector{Int64})
+    offspring = zero(a)
+    for i in a
+        offspring[i] += 1
+    end
+    return offspring
+end
+
+## FILTERING WITH ANCESTRY #################################################################
+
+mutable struct AncestryContainer{T,WT<:Real} <: AbstractParticleContainer{T}
+    tree::ParticleTree{T}
+    log_weights::Vector{WT}
+
+    function AncestryContainer(
+        initial_states::Vector{T}, log_weights::Vector{WT}, C::Int64=1
+    ) where {T,WT<:Real}
+        N = length(log_weights)
+        M = floor(C * N * log(N))
+        tree = ParticleTree(initial_states, Int64(M))
+        return new{T,WT}(tree, log_weights)
+    end
+end
+
+function Base.collect(ac::AncestryContainer)
+    return getindex(ac.tree.states, ac.tree.leaves)
+end
+
+function Base.getindex(ac::AncestryContainer, a::AbstractVector{Int64})
+    return getindex(ac.tree.states, getindex(ac.tree.leaves, a))
+end
+
+function reset_weights!(ac::AncestryContainer{T,WT}) where {T,WT<:Real}
+    fill!(ac.log_weights, zero(WT))
+    return ac.log_weights
+end
+
+function store!(ac::AncestryContainer, new_states, idx)
+    prune!(ac.tree, get_offspring(idx))
+    insert!(ac.tree, new_states, idx)
+    return ac
+end
+
+# start at each leaf and retrace it's steps to the root node
+function get_ancestry(tree::ParticleTree{T}) where {T}
+    paths = Vector{Vector{T}}(undef, length(tree.leaves))
+    @inbounds for (k, i) in enumerate(tree.leaves)
+        j = tree.parents[i]
+        xi = tree.states[i]
+
+        xs = [xi]
+        while j > 0
+            push!(xs, tree.states[j])
+            j = tree.parents[j]
+        end
+        paths[k] = reverse(xs)
+    end
+    return paths
+end

examples/particle-mcmc/resamplers.jl

@@ -0,0 +1,84 @@
+using Random
+using Distributions
+
+function multinomial_resampling(
+    rng::AbstractRNG, weights::AbstractVector{WT}, n::Int64=length(weights); kwargs...
+) where {WT<:Real}
+    return rand(rng, Distributions.Categorical(weights), n)
+end
+
+function systematic_resampling(
+    rng::AbstractRNG, weights::AbstractVector{WT}, n::Int64=length(weights); kwargs...
+) where {WT<:Real}
+    # pre-calculations
+    @inbounds v = n * weights[1]
+    u = oftype(v, rand(rng))
+
+    # initialize sampling algorithm
+    a = Vector{Int64}(undef, n)
+    idx = 1
+
+    @inbounds for i in 1:n
+        while v < u
+            idx += 1
+            v += n * weights[idx]
+        end
+        a[i] = idx
+        u += one(u)
+    end
+
+    return a
+end
+
+# TODO: this should be done in the log domain and also parallelized
+function metropolis_resampling(
+    rng::AbstractRNG,
+    weights::AbstractVector{WT},
+    n::Int64=length(weights);
+    ε::Float64=0.01,
+    kwargs...,
+) where {WT<:Real}
+    # pre-calculations
+    β = mean(weights)
+    bins = Int64(cld(log(ε), log(1 - β)))
+
+    # initialize the algorithm
+    a = Vector{Int64}(undef, n)
+
+    @inbounds for i in 1:n
+        k = i
+        for _ in 1:bins
+            j = rand(rng, 1:n)
+            v = weights[j] / weights[k]
+            if rand(rng) ≤ v
+                k = j
+            end
+        end
+        a[i] = k
+    end
+
+    return a
+end
+
+# TODO: this should be done in the log domain and also parallelized
+function rejection_resampling(
+    rng::AbstractRNG, weights::AbstractVector{WT}, n::Int64=length(weights); kwargs...
+) where {WT<:Real}
+    # pre-calculations
+    max_weight = maximum(weights)
+
+    # initialize the algorithm
+    a = Vector{Int64}(undef, n)
+
+    @inbounds for i in 1:n
+        j = i
+        u = rand(rng)
+        while u > weights[j] / max_weight
+            j = rand(1:n)
+            u = rand(rng)
+        end
+        a[i] = j
+    end
+
+    return a
+end

examples/particle-mcmc/script.jl

@@ -0,0 +1,97 @@
+using AdvancedMH
+using CairoMakie
+using StatsBase: weights, mean
+
+include("particles.jl")
+include("resamplers.jl")
+include("simple-filters.jl")
+
+## FILTERING DEMONSTRATION #################################################################
+
+# use a local level trend model
+function simulation_model(σx²::T, σy²::T) where {T<:Real}
+    init = Gaussian(zeros(T, 2), PDMat(diagm(ones(T, 2))))
+    dyn = LinearGaussianLatentDynamics(T[1 1; 0 1], T[0; 0], [σx² 0; 0 0], init)
+    obs = LinearGaussianObservationProcess(T[1 0], [σy²;;])
+    return StateSpaceModel(dyn, obs)
+end
+
+true_params = randexp(Float32, 2);
+true_model = simulation_model(true_params...);
+
+# simulate data
+rng = MersenneTwister(1234);
+_, _, data = sample(rng, true_model, 150);
+
+# test the adaptive resampling procedure
+states, llbf = sample(rng, true_model, data, BF(2048, 0.5); store_ancestry=true);
+
+# plot the smoothed states to validate the algorithm
+smoothed_trend = begin
+    fig = Figure(; size=(1200, 400))
+    ax1 = Axis(fig[1, 1])
+    ax2 = Axis(fig[1, 2])
+
+    # this is gross but it works fro visualization purposes
+    all_paths = map(x -> hcat(x...), get_ancestry(states.tree))
+    mean_paths = mean(all_paths, weights(softmax(states.log_weights)))
+    n_paths = length(all_paths)
+
+    # plot smoothed states in black and observed data in red
+    lines!(ax1, mean_paths[1, :]; color=:black)
+    lines!(ax1, vcat(0, data...); color=:red, linestyle=:dash)
+
+    # plot ancestry tree in graded black and data in red
+    lines!.(ax2, getindex.(all_paths, 1, :), color=(:black, maximum([2 / n_paths, 1e-2])))
+    lines!(ax2, vcat(0, data...); color=:red, linestyle=:dash)
+
+    fig
+end
+
+## PARTICLE MCMC ###########################################################################
+
+# consider a default Gamma prior with Float32s
+prior_dist = product_distribution(Gamma(1.0f0), Gamma(1.0f0));
+
+# basic RWMH ala AdvancedMH
+function density(θ::Vector{T}) where {T<:Real}
+    if insupport(prior_dist, θ)
+        # _, ll = sample(rng, simulation_model(θ...), data, BF(512))
+        _, ll = sample(rng, simulation_model(θ...), data, KF())
+        return ll + logpdf(prior_dist, θ)
+    else
+        return -Inf
+    end
+end
+
+pmmh = RWMH(MvNormal(zeros(Float32, 2), (0.01f0) * I));
+model = DensityModel(density);
+
+# works with AdvancedMH out of the box
+chains = sample(model, pmmh, 50_000);
+burn_in = 1_000;
+
+# plot the posteriors
+hist_plots = begin
+    param_post = hcat(getproperty.(chains[burn_in:end], :params)...)
+    fig = Figure(; size=(1200, 400))
+
+    for i in 1:2
+        # plot the posteriors with burn-in
+        hist(
+            fig[1, i],
+            param_post[i, :];
+            color=(:black, 0.4),
+            strokewidth=1,
+            normalization=:pdf,
+        )
+
+        # plot the true values
+        vlines!(fig[1, i], true_params[i]; color=:red, linestyle=:dash, linewidth=3)
+    end
+
+    fig
+end
+
+# this is useful for SMC algorithms like SMC² or density tempered SMC
+acc_ratio = mean(getproperty.(chains, :accepted))