Draft DA script

examples/particle-mcmc/Project.toml

@@ -0,0 +1,16 @@
+[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"
+GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
+GaussianDistributions = "43dcc890-d446-5863-8d1a-14597580bb8d"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+SSMProblems = "26aad666-b158-4e64-9d35-0e672562fa48"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
+UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"

examples/particle-mcmc/lorenz.jl

@@ -0,0 +1,93 @@
+using Distributions
+using Random
+using SSMProblems
+using UnPack
+using OrdinaryDiffEq
+using LinearAlgebra
+using PDMats
+using GLMakie
+
+include("particles.jl")
+include("resamplers.jl")
+include("simple-filters.jl")
+
+Base.@kwdef struct Parameters{T<:Real}
+    β::T = 8 / 3
+    ρ::T = 28.0
+    σ::T = 10.0
+    ν::T = 1.0 # Obs noise variance
+    dt::T = 0.025 # Time step
+end
+
+function lorenz!(du, u, p::Parameters, t)
+    @unpack β, ρ, σ = p
+    du[1] = σ * (u[2] - u[1])
+    du[2] = u[1] * (ρ - u[3]) - u[2]
+    return du[3] = u[1] * u[2] - β * u[3]
+end
+
+struct LatentNoiseProcess{T} <: LatentDynamics{Vector{T}}
+    σ::AbstractPDMat{T}
+    dt::T
+    integrator
+end
+
+struct ObservationNoiseProcess{T} <: ObservationProcess{Vector{T}}
+    σ::AbstractPDMat{T}
+end
+
+function SSMProblems.distribution(dyn::LatentNoiseProcess, step::Integer, prev_state, extra)
+    reinit!(dyn.integrator, prev_state)
+    step!(dyn.integrator, dyn.dt, true)
+    return MvNormal(dyn.integrator.u, dyn.σ)
+end
+
+function SSMProblems.distribution(dyn::LatentNoiseProcess, extra)
+    return MvNormal([1; 0; 0], dyn.σ)
+end
+
+function SSMProblems.distribution(obs::ObservationNoiseProcess, step::Integer, state, extra)
+    return MvNormal(state, obs.σ * I)
+end
+
+# Simulate some data
+u0 = [1.0; 0.0; 0.0]
+params = Parameters()
+
+dt = 0.025
+N = 100
+Np = 512
+tspan = (0.0, dt * N)
+
+rng = MersenneTwister()
+
+prob = ODEProblem(lorenz!, u0, tspan, params)
+alg = Tsit5()
+integrator = init(prob, Tsit5(); dt=dt, adaptive=false)
+sol = solve(prob, alg; dt=dt, adaptive=false)
+
+# SSM Noise Model
+dyn = LatentNoiseProcess(ScalMat(3, params.dt), params.dt, integrator)
+obs = ObservationNoiseProcess(ScalMat(3, params.ν))
+model = StateSpaceModel(dyn, obs)
+x0, x, y = sample(rng, model, N)
+
+filter = BF(Np; threshold=1.0, resampler=Systematic());
+sparse_ancestry = AncestorCallback(eltype(model.dyn), filter.N, 1.0);
+tree, llbf = sample(rng, model, filter, y; callback=sparse_ancestry);
+lineage = get_ancestry(sparse_ancestry.tree)
+
+# Fancy 3D plot
+# fig = Figure()
+# lines(fig[1, 1], hcat(x0, x...))
+# for (i, path) in enumerate(lineage)
+# 	lines!(fig[1, 1], hcat(path...), color=:black)
+# end
+
+fig = Figure()
+for i in eachindex(first(x))
+    lines(fig[i, 1], hcat(x0, x...)[i, :])
+    for path in lineage
+        lines!(fig[i, 1], hcat(path...)[i, :]; color=:black, alpha=0.1)
+    end
+end

examples/particle-mcmc/particles.jl

@@ -0,0 +1,176 @@
+using DataStructures: Stack
+using StatsBase
+
+## PARTICLES ###############################################################################
+
+mutable struct ParticleContainer{T,WT<:Real}
+    filtered::Vector{T}
+    proposed::Vector{T}
+    ancestors::Vector{Int64}
+    log_weights::Vector{WT}
+
+    function ParticleContainer(
+        initial_states::Vector{T}, log_weights::Vector{WT}
+    ) where {T,WT<:Real}
+        return new{T,WT}(
+            initial_states, similar(initial_states), eachindex(log_weights), log_weights
+        )
+    end
+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 reset_weights!(pc::ParticleContainer{T,WT}) where {T,WT<:Real}
+    fill!(pc.log_weights, zero(WT))
+    return pc.log_weights
+end
+
+function StatsBase.weights(pc::ParticleContainer)
+    return softmax(pc.log_weights)
+end
+
+## SPARSE 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
+
+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
+
+## ANCESTOR STORAGE CALLBACK ###############################################################
+
+mutable struct AncestorCallback
+    tree::ParticleTree
+
+    function AncestorCallback(::Type{T}, N::Integer, C::Real=1.0) where {T}
+        M = floor(Int64, C * N * log(N))
+        nodes = Vector{T}(undef, N)
+        return new(ParticleTree(nodes, M))
+    end
+end
+
+function (c::AncestorCallback)(model, filter, step, states, data; kwargs...)
+    if step == 1
+        # this may be incorrect, but it is functional
+        @inbounds c.tree.states[1:(filter.N)] = deepcopy(states.filtered)
+    end
+    prune!(c.tree, get_offspring(states.ancestors))
+    insert!(c.tree, states.filtered, states.ancestors)
+    return nothing
+end
+
+mutable struct ResamplerCallback
+    tree::ParticleTree
+
+    function ResamplerCallback(N::Integer, C::Real=1.0)
+        M = floor(Int64, C * N * log(N))
+        nodes = collect(1:N)
+        return new(ParticleTree(nodes, M))
+    end
+end
+
+function (c::ResamplerCallback)(model, filter, step, states, data; kwargs...)
+    if step != 1
+        prune!(c.tree, get_offspring(states.ancestors))
+        insert!(c.tree, collect(1:(filter.N)), states.ancestors)
+    end
+    return nothing
+end

examples/particle-mcmc/resamplers.jl

@@ -0,0 +1,102 @@
+using Random
+using Distributions
+
+abstract type AbstractResampler end
+
+## DOUBLE PRECISION STABLE ALGORITHMS ######################################################
+
+struct Multinomial <: AbstractResampler end
+
+function resample(
+    rng::AbstractRNG, ::Multinomial, weights::AbstractVector{WT}, n::Int64=length(weights)
+) where {WT<:Real}
+    return rand(rng, Distributions.Categorical(weights), n)
+end
+
+struct Systematic <: AbstractResampler end
+
+function resample(
+    rng::AbstractRNG, ::Systematic, weights::AbstractVector{WT}, n::Int64=length(weights)
+) 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
+
+## SINGLE PRECISION STABLE ALGORITHMS ######################################################
+
+struct Metropolis <: AbstractResampler
+    ε::Float64
+    function Metropolis(ε::Float64=0.01)
+        return new(ε)
+    end
+end
+
+# TODO: this should be done in the log domain and also parallelized
+function resample(
+    rng::AbstractRNG,
+    resampler::Metropolis,
+    weights::AbstractVector{WT},
+    n::Int64=length(weights);
+) where {WT<:Real}
+    # pre-calculations
+    β = mean(weights)
+    B = Int64(cld(log(resampler.ε), log(1 - β)))
+
+    # initialize the algorithm
+    a = Vector{Int64}(undef, n)
+
+    @inbounds for i in 1:n
+        k = i
+        for _ in 1:B
+            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
+
+struct Rejection <: AbstractResampler end
+
+# TODO: this should be done in the log domain and also parallelized
+function resample(
+    rng::AbstractRNG, ::Rejection, weights::AbstractVector{WT}, n::Int64=length(weights)
+) 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(rng, 1:n)
+            u = rand(rng)
+        end
+        a[i] = j
+    end
+
+    return a
+end