Make HierarchicalSSM compatible with the BF

GeneralisedFilters/src/algorithms/bootstrap.jl

@@ -104,3 +104,19 @@ function update(
 
     return filtered, ll_increment
 end
+
+# Application of bootstrap filter to hierarchical models
+function filter(
+    rng::AbstractRNG,
+    model::HierarchicalSSM,
+    alg::BootstrapFilter,
+    observations::AbstractVector;
+    ref_state::Union{Nothing,AbstractVector}=nothing,
+    kwargs...,
+)
+    ssm = StateSpaceModel(
+        HierarchicalDynamics(model.outer_dyn, model.inner_model.dyn),
+        HierarchicalObservations(model.inner_model.obs),
+    )
+    return filter(rng, ssm, alg, observations; ref_state=ref_state, kwargs...)
+end

GeneralisedFilters/src/models/hierarchical.jl

@@ -48,3 +48,55 @@ function AbstractMCMC.sample(
 
     return x0, z0, xs, zs, ys
 end
+
+## Methods to make HierarchicalSSM compatible with the bootstrap filter
+struct HierarchicalDynamics{T<:Real,ET,D1<:LatentDynamics{T},D2<:LatentDynamics{T}} <:
+       LatentDynamics{T,ET}
+    outer_dyn::D1
+    inner_dyn::D2
+    function HierarchicalDynamics(
+        outer_dyn::D1, inner_dyn::D2
+    ) where {D1<:LatentDynamics,D2<:LatentDynamics}
+        ET = RaoBlackwellisedParticle{eltype(outer_dyn),eltype(inner_dyn)}
+        T = SSMProblems.arithmetic_type(outer_dyn)
+        return new{T,ET,D1,D2}(outer_dyn, inner_dyn)
+    end
+end
+
+function SSMProblems.simulate(rng::AbstractRNG, dyn::HierarchicalDynamics; kwargs...)
+    outer_dyn, inner_dyn = dyn.outer_dyn, dyn.inner_dyn
+    x0 = simulate(rng, outer_dyn; kwargs...)
+    z0 = simulate(rng, inner_dyn; new_outer=x0, kwargs...)
+    return RaoBlackwellisedParticle(x0, z0)
+end
+
+function SSMProblems.simulate(
+    rng::AbstractRNG,
+    dyn::HierarchicalDynamics,
+    step::Integer,
+    prev_state::RaoBlackwellisedParticle;
+    kwargs...,
+)
+    outer_dyn, inner_dyn = dyn.outer_dyn, dyn.inner_dyn
+    x = simulate(rng, outer_dyn, step, prev_state.x; kwargs...)
+    z = simulate(
+        rng, inner_dyn, step, prev_state.z; prev_outer=prev_state.x, new_outer=x, kwargs...
+    )
+    return RaoBlackwellisedParticle(x, z)
+end
+
+struct HierarchicalObservations{T<:Real,ET,OP<:ObservationProcess{T}} <:
+       ObservationProcess{T,ET}
+    obs::OP
+    function HierarchicalObservations(obs::OP) where {OP<:ObservationProcess}
+        T = SSMProblems.arithmetic_type(obs)
+        ET = eltype(obs)
+        return new{T,ET,OP}(obs)
+    end
+end
+
+function SSMProblems.distribution(
+    obs::HierarchicalObservations, step::Integer, state::RaoBlackwellisedParticle; kwargs...
+)
+    return distribution(obs.obs, step, state.z; new_outer=state.x, kwargs...)
+end

GeneralisedFilters/test/runtests.jl

@@ -295,6 +295,46 @@ end
     paths = GeneralisedFilters.get_ancestry(tree)
 end
 
+@testitem "BF on hierarchical model test" begin
+    using LogExpFunctions: softmax
+    using StableRNGs
+    using StatsBase
+
+    SEED = 1234
+    D_outer = 1
+    D_inner = 1
+    D_obs = 1
+    T = 5
+    N_particles = 10^4
+
+    rng = StableRNG(SEED)
+
+    full_model, hier_model = GeneralisedFilters.GFTest.create_dummy_linear_gaussian_model(
+        rng, D_outer, D_inner, D_obs
+    )
+    _, _, ys = sample(rng, full_model, T)
+
+    # Ground truth Kalman filtering
+    kf_states, kf_ll = GeneralisedFilters.filter(rng, full_model, KalmanFilter(), ys)
+
+    # Rao-Blackwellised particle filtering
+    bf = BF(N_particles)
+    states, ll = GeneralisedFilters.filter(rng, hier_model, bf, ys)
+
+    # Extract final filtered states
+    xs = map(p -> getproperty(p, :x), states.particles)
+    zs = map(p -> getproperty(p, :z), states.particles)
+    log_ws = states.log_weights
+
+    @test kf_ll ≈ ll rtol = 1e-2
+
+    weights = Weights(softmax(log_ws))
+
+    # Higher tolerance for outer state since variance is higher
+    @test first(kf_states.μ) ≈ sum(only.(xs) .* weights) rtol = 1e-1
+    @test last(kf_states.μ) ≈ sum(only.(zs) .* weights) rtol = 1e-1
+end
+
 @testitem "Dense ancestry test" begin
     using GeneralisedFilters
     using StableRNGs