Generalise Container Types (#151)

* Add handling for PD errors

* Generalise types for containers and add ReferenceTrajectory type

* Formatter suggestions

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: THargreaves

GeneralisedFilters/Project.toml

@@ -16,7 +16,6 @@ LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
-OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SSMProblems = "26aad666-b158-4e64-9d35-0e672562fa48"
@@ -56,7 +55,6 @@ LogDensityProblemsAD = "1.13.1"
 LogExpFunctions = "0.3"
 MCMCChains = "7.7.0"
 Mooncake = "0.5.13"
-OffsetArrays = "1.14.1"
 PDMats = "0.11.35"
 SSMProblems = "0.6"
 StaticArrays = "1.9.17"

GeneralisedFilters/ext/CUDAExt.jl

@@ -19,7 +19,7 @@ using GeneralisedFilters:
     PostUpdateCallback,
     num_particles
 
-using GeneralisedFilters.OffsetArrays: OffsetVector
+using GeneralisedFilters: ReferenceTrajectory
 
 using AcceleratedKernels: searchsortedfirst, foreachindex
 using CUDA
@@ -226,29 +226,37 @@ function expand!(tree::ParallelParticleTree{T}) where {T}
     return tree
 end
 
-# Get ancestry of all particles
+# Get ancestry of all particles. The parallel tree stores initial and subsequent states
+# in a single `states` buffer (homogeneous type), so the returned trajectories have
+# T0 == T.
 function get_ancestry(tree::ParallelParticleTree{ST}, T::Integer) where {ST}
-    paths = OffsetVector(Vector{ST}(undef, T + 1), -1)
+    buf = Vector{Vector{ST}}(undef, T + 1)
     parents = copy(tree.leaves)
-    for t in T:-1:1
-        paths[t] = tree.states[parents]
+    for t in (T + 1):-1:2
+        buf[t] = Vector(tree.states[parents])
         gather!(parents, tree.parents, parents)
     end
-    return paths
+    buf[1] = Vector(tree.states[parents])
+    # Each leaf's trajectory: x0 = buf[1][k], xs = [buf[2][k], ..., buf[T+1][k]]
+    return [
+        ReferenceTrajectory(buf[1][k], [buf[t + 1][k] for t in 1:T]) for
+        k in eachindex(tree.leaves)
+    ]
 end
 
 # Get ancestry of a single particle
 function get_ancestry(
     container::ParallelParticleTree{ST}, i::Integer, T::Integer
 ) where {ST}
-    path = OffsetVector(Vector{ST}(undef, T + 1), -1)
+    xs = Vector{ST}(undef, T)
     CUDA.@allowscalar begin
         ancestor_index = container.leaves[i]
-        for t in T:-1:0
-            path[t] = container.states[ancestor_index]
+        for t in T:-1:1
+            xs[t] = container.states[ancestor_index]
             ancestor_index = container.parents[ancestor_index]
         end
-        return path
+        x0 = container.states[ancestor_index]
+        return ReferenceTrajectory(x0, xs)
     end
 end
 

GeneralisedFilters/ext/TuringExt.jl

@@ -20,7 +20,6 @@ using DynamicPPL: DynamicPPL
 using LinearAlgebra: PosDefException
 using LogDensityProblems: LogDensityProblems
 using MCMCChains: MCMCChains
-using OffsetArrays
 using Random: AbstractRNG
 using SSMProblems
 using Turing: Turing

GeneralisedFilters/src/GeneralisedFilters.jl

@@ -5,16 +5,15 @@ using ADTypes: ADTypes
 import Distributions: MvNormal, params
 import Random: AbstractRNG, default_rng, rand
 import SSMProblems: prior, dyn, obs
-using OffsetArrays
 using SSMProblems
 using StatsBase
 using DifferentiationInterface
 
 const DI = DifferentiationInterface
 
 # Filtering utilities
-include("callbacks.jl")
 include("containers.jl")
+include("callbacks.jl")
 include("resamplers.jl")
 
 ## FILTERING BASE ##########################################################################

GeneralisedFilters/src/algorithms/csmc.jl

@@ -72,8 +72,8 @@ ConditionalSMC(pf) = ConditionalSMC(pf, NoRefreshment())
 
 State of a conditional SMC sampler, containing the current reference trajectory.
 
-The trajectory is an `OffsetVector` indexed from 0 (matching the prior at time 0).
-For RBPF, the trajectory contains `RBState` objects (outer state + inner filtering
+The trajectory is a [`ReferenceTrajectory`](@ref) indexed from 0 (matching the prior at
+time 0). For RBPF, the trajectory contains `RBState` objects (outer state + inner filtering
 distribution).
 """
 struct CSMCState{TT}
@@ -100,10 +100,11 @@ end
 # CSMCState stores full trajectories (RBState for RBPF). The filter/initialise/move
 # functions expect ref_state to contain only outer states for RBPF. _make_ref_state
 # handles this conversion.
-# TODO: is this actually needed?
 _make_ref_state(::Nothing) = nothing
-_make_ref_state(traj) = traj
-_make_ref_state(traj::OffsetVector{<:RBState}) = map(s -> s.x, traj)
+_make_ref_state(traj::ReferenceTrajectory) = traj
+function _make_ref_state(traj::ReferenceTrajectory{<:RBState})
+    return map(s -> s.x, traj)
+end
 
 ## TRAJECTORY SAMPLING #####################################################################
 
@@ -119,16 +120,34 @@ function _sample_trajectory(
     rng::AbstractRNG, tree::ParticleTree, state::ParticleDistribution
 )
     ws = get_weights(state)
-    path = rand(rng, tree, ws)
-    return OffsetVector(path, -1)
+    return rand(rng, tree, ws)
 end
 
-## PARTICLE TREE HELPERS ###################################################################
+## PARTICLE TREE / CONTAINER HELPERS ######################################################
+
+# Capacity heuristic from Jacob, Murray & Rubenthaler (2015)
+_tree_capacity(N::Integer) = max(N, floor(Int64, N * log(N)))
+
+# Construct a ParticleTree using both the time-0 and time-1 particle distributions so
+# that the subsequent-state type `T` is inferred from the time-1 states (which may
+# differ from the type of the initial states in Rao-Blackwellised settings).
+function _init_tree(init_state::ParticleDistribution, state::ParticleDistribution)
+    initial_states = map(p -> p.state, init_state.particles)
+    states_t1 = map(p -> p.state, state.particles)
+    ancestors_t1 = map(p -> p.ancestor, state.particles)
+    return ParticleTree(
+        initial_states, states_t1, ancestors_t1, _tree_capacity(length(initial_states))
+    )
+end
 
-function _init_tree(state::ParticleDistribution)
-    states = map(p -> p.state, state.particles)
-    N = length(states)
-    return ParticleTree(states, max(N, floor(Int64, N * log(N))))
+function _init_container(init_state::ParticleDistribution, state::ParticleDistribution)
+    initial_states = map(p -> p.state, init_state.particles)
+    return DenseParticleContainer(
+        initial_states,
+        map(p -> p.state, state.particles),
+        Float64.(log_weights(state)),
+        map(p -> p.ancestor, state.particles),
+    )
 end
 
 function _update_tree!(tree::ParticleTree, state::ParticleDistribution)
@@ -140,6 +159,17 @@ function _update_tree!(tree::ParticleTree, state::ParticleDistribution)
     return tree
 end
 
+function _update_container!(c::DenseParticleContainer, state::ParticleDistribution)
+    particles = state.particles
+    push!(
+        c,
+        map(p -> p.state, particles),
+        Float64.(log_weights(state)),
+        map(p -> p.ancestor, particles),
+    )
+    return c
+end
+
 ## BACKWARD PREDICTIVE LIKELIHOODS #########################################################
 
 # Default: no backward likelihoods needed (regular PF, or first iteration)
@@ -217,11 +247,9 @@ function _csmc_sample(
     K = length(observations)
     ref_state = _make_ref_state(ref_traj)
 
-    state = initialise(rng, prior(model), pf; ref_state)
-    tree = _init_tree(state)
-
-    state, ll = step(rng, model, pf, 1, state, observations[1]; ref_state)
-    _update_tree!(tree, state)
+    init_state = initialise(rng, prior(model), pf; ref_state)
+    state, ll = step(rng, model, pf, 1, init_state, observations[1]; ref_state)
+    tree = _init_tree(init_state, state)
 
     for t in 2:K
         state, ll_inc = step(rng, model, pf, t, state, observations[t]; ref_state)
@@ -247,19 +275,17 @@ function _csmc_sample(
     # Backward predictive likelihoods (only non-nothing for RBPF)
     back_liks = _compute_backward_likelihoods(rng, model, pf, observations, ref_state)
 
-    state = initialise(rng, prior(model), pf; ref_state)
-    tree = _init_tree(state)
+    init_state = initialise(rng, prior(model), pf; ref_state)
 
-    ll = zero(eltype(log_weights(state)))
-    for t in 1:K
-        # Ancestor sampling for the reference particle
+    # Perform one CSMC-AS step on the current state
+    function _csmc_as_step(state, t)
+        ancestor_idx = 0
         if !isnothing(ref_state)
             ref_as = _build_ancestor_ref(ref_state, back_liks, t)
             as_weights = map(state.particles) do particle
                 ancestor_weight(particle, dyn(model), pf, t, ref_as)
             end
             ancestor_idx = StatsBase.sample(rng, StatsBase.Weights(softmax(as_weights)))
-            
         end
 
         state = resample(rng, resampler(pf), state; ref_state)
@@ -270,9 +296,15 @@ function _csmc_sample(
             )
         end
 
-        state, ll_inc = move(rng, model, pf, t, state, observations[t]; ref_state)
-        ll += ll_inc
+        return move(rng, model, pf, t, state, observations[t]; ref_state)
+    end
+
+    state, ll = _csmc_as_step(init_state, 1)
+    tree = _init_tree(init_state, state)
 
+    for t in 2:K
+        state, ll_inc = _csmc_as_step(state, t)
+        ll += ll_inc
         _update_tree!(tree, state)
     end
 
@@ -366,55 +398,51 @@ function _csmc_sample(
 )
     pf = csmc.pf
     K = length(observations)
+    N = num_particles(pf)
     ref_state = _make_ref_state(ref_traj)
 
-    # Forward filtering pass, storing particles at each timestep
+    # Forward filtering pass: store full history in a DenseParticleContainer.
     init_state = initialise(rng, prior(model), pf; ref_state)
-    init_particles = copy(init_state.particles)
-
     state, ll = step(rng, model, pf, 1, init_state, observations[1]; ref_state)
-    particle_history = Vector{typeof(state.particles)}(undef, K)
-    particle_history[1] = copy(state.particles)
+    container = _init_container(init_state, state)
 
     for t in 2:K
         state, ll_inc = step(rng, model, pf, t, state, observations[t]; ref_state)
         ll += ll_inc
-        particle_history[t] = copy(state.particles)
+        _update_container!(container, state)
     end
 
     # Backward simulation pass
-    ws = get_weights(state)
-    idx = StatsBase.sample(rng, StatsBase.Weights(ws))
-    sampled_state = particle_history[K][idx].state
+    idx = StatsBase.sample(rng, StatsBase.Weights(get_weights(state)))
+    sampled_state = container.states[K][idx]
 
     back_lik = _bs_init_back_lik(rng, model, pf, observations, K, sampled_state)
 
-    ST = typeof(sampled_state)
-    trajectory = OffsetVector(Vector{ST}(undef, K + 1), -1)
-    trajectory[K] = sampled_state
+    xs = Vector{typeof(sampled_state)}(undef, K)
+    xs[K] = sampled_state
 
     for t in (K - 1):-1:1
-        ref_next = _build_bs_ref(trajectory[t + 1], back_lik)
-        backward_ws = map(particle_history[t]) do particle
-            ancestor_weight(particle, dyn(model), pf, t + 1, ref_next)
+        ref_next = _build_bs_ref(xs[t + 1], back_lik)
+        backward_ws = map(1:N) do i
+            ancestor_weight(Particle(container, t, i), dyn(model), pf, t + 1, ref_next)
         end
         idx = StatsBase.sample(rng, StatsBase.Weights(softmax(backward_ws)))
-        trajectory[t] = particle_history[t][idx].state
+        xs[t] = container.states[t][idx]
 
         back_lik = _bs_step_back_lik(
-            rng, model, pf, t, back_lik, observations, trajectory[t], trajectory[t + 1]
+            rng, model, pf, t, back_lik, observations, xs[t], xs[t + 1]
         )
     end
 
-    # Time 0: backward step from t=1 to initial particles
-    ref_at_1 = _build_bs_ref(trajectory[1], back_lik)
-    backward_ws = map(init_particles) do particle
+    # Time 0: backward step from t=1 to initial particles.
+    ref_at_1 = _build_bs_ref(xs[1], back_lik)
+    backward_ws = map(init_state.particles) do particle
         ancestor_weight(particle, dyn(model), pf, 1, ref_at_1)
     end
     idx = StatsBase.sample(rng, StatsBase.Weights(softmax(backward_ws)))
-    trajectory[0] = init_particles[idx].state
+    x0 = container.initial_states[idx]
 
-    return trajectory, ll
+    return ReferenceTrajectory(x0, xs), ll
 end
 
 ## ABSTRACTMCMC INTERFACE ##################################################################