Weigghts:

Author: FredericWantiez

src/algorithms/apf.jl

@@ -1,7 +1,7 @@
 export AuxiliaryParticleFilter, APF
 
-mutable struct AuxiliaryParticleFilter{RS<:AbstractConditionalResampler} <: AbstractFilter
-    N::Integer
+mutable struct AuxiliaryParticleFilter{N,RS<:AbstractConditionalResampler} <:
+               AbstractParticleFilter{N}
     resampler::RS
     aux::Vector # Auxiliary weights
 end
@@ -10,20 +10,22 @@ function AuxiliaryParticleFilter(
     N::Integer; threshold::Real=0.0, resampler::AbstractResampler=Systematic()
 )
     conditional_resampler = ESSResampler(threshold, resampler)
-    return AuxiliaryParticleFilter(N, conditional_resampler, zeros(N))
+    return AuxiliaryParticleFilter{N,typeof(conditional_resampler)}(
+        conditional_resampler, zeros(N)
+    )
 end
 
 const APF = AuxiliaryParticleFilter
 
 function initialise(
     rng::AbstractRNG,
     model::StateSpaceModel{T},
-    filter::AuxiliaryParticleFilter;
+    filter::AuxiliaryParticleFilter{N};
     ref_state::Union{Nothing,AbstractVector}=nothing,
     kwargs...,
-) where {T}
-    initial_states = map(x -> SSMProblems.simulate(rng, model.dyn; kwargs...), 1:(filter.N))
-    initial_weights = fill(-log(T(filter.N)), filter.N)
+) where {T,N}
+    initial_states = map(x -> SSMProblems.simulate(rng, model.dyn; kwargs...), 1:N)
+    initial_weights = fill(-log(T(N)), N)
 
     return update_ref!(
         ParticleContainer(initial_states, initial_weights), ref_state, filter
@@ -75,12 +77,12 @@ end
 
 function update(
     model::StateSpaceModel{T},
-    filter::AuxiliaryParticleFilter,
+    filter::AuxiliaryParticleFilter{N},
     step::Integer,
     states::ParticleContainer,
     observation;
     kwargs...,
-) where {T}
+) where {T,N}
     @debug "step $step"
     log_increments = map(
         x -> SSMProblems.logdensity(model.obs, step, x, observation; kwargs...),
@@ -90,7 +92,7 @@ function update(
     states.filtered.log_weights = states.proposed.log_weights + log_increments
     states.filtered.particles = states.proposed.particles
 
-    return (states, logsumexp(log_increments) - log(T(filter.N)))
+    return (states, logsumexp(log_increments) - log(T(N)))
 end
 
 function step(

src/algorithms/bootstrap.jl

@@ -71,7 +71,7 @@ function update(
 end
 
 function reset_weights!(
-    state::ParticleState{T,WT}, idxs, filter::BootstrapFilter{N}
+    state::ParticleState{T,WT}, idxs, ::BootstrapFilter{N}
 ) where {T,WT<:Real,N}
     fill!(state.log_weights, -log(WT(N)))
     return state

src/algorithms/ffbs.jl

@@ -25,7 +25,7 @@ function (callback::WeightedParticleRecorderCallback)(
     return nothing
 end
 
-function smooth(
+function sample(
     rng::Random.AbstractRNG,
     model::StateSpaceModel{T,LDT},
     alg::FFBS{<:BootstrapFilter{N}},
@@ -40,21 +40,34 @@ function smooth(
     )
 
     particles, _ = filter(rng, model, alg.filter, obs; callback=recorder, kwargs...)
+
+    # Backward sampling - exact
     idx_ref = rand(rng, Categorical(weights(particles.filtered)), M)
     trajectories = Array{eltype(model.dyn)}(undef, n_timestep, M)
 
     trajectories[end, :] = particles.filtered[idx_ref]
     for step in (n_timestep - 1):-1:1
         for j in 1:M
-            transitions = map(
-                x ->
-                    SSMProblems.logdensity(model.dyn, step, x, trajectories[step+1]; kwargs...),
+            backward_weights = backward(
+                model::StateSpaceModel,
+                step,
+                trajectories[step + 1],
                 recorder.particles[step, :],
+                recorder.log_weights[step, :];
+                kwargs...,
             )
-            backward_weights = recorder.log_weights[step, :] + transitions
             ancestor = rand(rng, Categorical(softmax(backward_weights)))
             trajectories[step, j] = recorder.particles[step, ancestor]
         end
     end
     return trajectories
 end
+
+function backward(
+    model::StateSpaceModel, step::Integer, state, particles::T, log_weights::WT; kwargs...
+) where {T,WT}
+    transitions = map(
+        x -> SSMProblems.logdensity(model.dyn, step, x, state; kwargs...), particles
+    )
+    return log_weights + transitions
+end