Restore Type Stability

GeneralisedFilters/src/GFTest/resamplers.jl

@@ -18,16 +18,22 @@ mutable struct AlternatingResampler <: GeneralisedFilters.AbstractConditionalRes
     end
 end
 
-function GeneralisedFilters.will_resample(alt_resampler::AlternatingResampler, state)
+function GeneralisedFilters.will_resample(
+    alt_resampler::AlternatingResampler, state, weights
+)
     return alt_resampler.resample_next
 end
 
 function GeneralisedFilters.resample(
     rng::AbstractRNG,
     alt_resampler::AlternatingResampler,
-    state;
+    state,
+    weights;
     ref_state::Union{Nothing,AbstractVector}=nothing,
+    kwargs...,
 )
     alt_resampler.resample_next = !alt_resampler.resample_next
-    return GeneralisedFilters.resample(rng, alt_resampler.resampler, state; ref_state)
+    return GeneralisedFilters.resample(
+        rng, alt_resampler.resampler, state, weights; ref_state, kwargs...
+    )
 end

GeneralisedFilters/src/algorithms/particles.jl

@@ -44,14 +44,7 @@ function initialise(
         initialise_particle(rng, prior, algo, ref; kwargs...)
     end
 
-    # Set equal weights: log_w = -log(N) so weights sum to 1
-    log_weight = -log(N)
-    for p in particles
-        p.log_w = log_weight
-    end
-
-    # Initialize with ll_baseline = 0.0
-    return ParticleDistribution(particles, 0.0)
+    return ParticleDistribution(particles, false)
 end
 
 function predict(
@@ -64,20 +57,18 @@ function predict(
     ref_state::Union{Nothing,AbstractVector}=nothing,
     kwargs...,
 )
-    particles = map(1:length(state.particles)) do i
+    particles = map(1:num_particles(algo)) do i
         particle = state.particles[i]
         ref = !isnothing(ref_state) && i == 1 ? ref_state[iter] : nothing
         predict_particle(rng, dyn, algo, iter, particle, observation, ref; kwargs...)
     end
-    state.particles = particles
 
     # Accumulate the baseline with LSE of weights after prediction (before update)
     # For plain PF/guided: ll_baseline is 0.0 on entry, becomes LSE_before
     # For APF with resample: ll_baseline already stores negative correction; add LSE_before
-    LSE_before = logsumexp(map(p -> p.log_w, state.particles))
-    state.ll_baseline += LSE_before
-
-    return state
+    return ParticleDistribution(
+        particles, logsumexp(log_weights(state)) + state.ll_baseline
+    )
 end
 
 function update(
@@ -91,10 +82,9 @@ function update(
     particles = map(state.particles) do particle
         update_particle(obs, algo, iter, particle, observation; kwargs...)
     end
-    state.particles = particles
-    ll_increment = marginalise!(state)
+    new_state, ll_increment = marginalise!(state, particles)
 
-    return state, ll_increment
+    return new_state, ll_increment
 end
 
 struct ParticleFilter{RS,PT} <: AbstractParticleFilter
@@ -119,38 +109,37 @@ function initialise_particle(
     rng::AbstractRNG, prior::StatePrior, algo::ParticleFilter, ref_state; kwargs...
 )
     x = sample_prior(rng, prior, algo, ref_state; kwargs...)
-    # TODO (RB):  determine the correct type for the log_w field or use a NoWeight type
-    return Particle(x, 0.0, 0)
+    return Particle(x, 0)
 end
 
 function predict_particle(
     rng::AbstractRNG,
     dyn::LatentDynamics,
     algo::ParticleFilter,
     iter::Integer,
-    particle::Particle,
+    particle::AbstractParticle,
     observation,
     ref_state;
     kwargs...,
 )
-    new_x, logw_inc = propogate(
+    new_x, log_increment = propogate(
         rng, dyn, algo, iter, particle.state, observation, ref_state; kwargs...
     )
-    return Particle(new_x, particle.log_w + logw_inc, particle.ancestor)
+    return Particle(new_x, log_weight(particle) + log_increment, particle.ancestor)
 end
 
 function update_particle(
     obs::ObservationProcess,
     ::ParticleFilter,
     iter::Integer,
-    particle::Particle,
+    particle::AbstractParticle,
     observation;
     kwargs...,
 )
     log_increment = SSMProblems.logdensity(
         obs, iter, particle.state, observation; kwargs...
     )
-    return Particle(particle.state, particle.log_w + log_increment, particle.ancestor)
+    return Particle(particle.state, log_weight(particle) + log_increment, particle.ancestor)
 end
 
 function step(
@@ -250,10 +239,8 @@ function propogate(
         ref_state
     end
 
-    # TODO: make this type consistent
-    # Will have to do a lazy zero or change propogate to accept a particle (in which case
-    # we'll need to construct a particle in the RBPF predict method)
-    return new_x, 0.0
+    # TODO: replace this with nothing (unweighted particle)
+    return new_x, 0
 end
 
 # TODO: I feel like we shouldn't need to do this conversion. It should be handled by dispatch
@@ -312,40 +299,8 @@ function step(
         predictive_loglik(obs(model), algo.pf, iter, p_star, observation; kwargs...)
     end
 
-    # Log normalizer for current weights
-    LSE_w = logsumexp(map(p -> p.log_w, state.particles))
-
-    # Incorporate lookahead weights into current weights
-    for (i, particle) in enumerate(state.particles)
-        particle.log_w += log_ξs[i]
-    end
-
-    # Compute lookahead evidence
-    LSE_lookahead = logsumexp(map(p -> p.log_w, state.particles)) - LSE_w
-
-    resample_flag = will_resample(resampler(algo), state)
-    if resample_flag
-        state = resample(rng, resampler(algo), state; ref_state)
-    else
-        # Not resampling: preserve ll_baseline and set ancestors to self
-        n = length(state.particles)
-        new_particles = similar(state.particles)
-        for i in 1:n
-            new_particles[i] = set_ancestor(state.particles[i], i)
-        end
-        state = ParticleDistribution(new_particles, state.ll_baseline)
-    end
-
-    # Compensate for lookahead weights
-    for particle in state.particles
-        particle.log_w -= log_ξs[particle.ancestor]
-    end
-
-    # Compute compensation log normalizer
-    if resample_flag
-        LSE_comp = logsumexp(map(p -> p.log_w, state.particles))
-        state.ll_baseline = -(LSE_lookahead + (LSE_comp - log(num_particles(algo))))
-    end
+    rs = AuxiliaryResampler(resampler(algo), log_ξs)
+    state = maybe_resample(rng, rs, state; ref_state)
 
     callback(model, algo, iter, state, observation, PostResample; kwargs...)
     return move(
@@ -386,11 +341,11 @@ function predictive_state(
     dyn::LatentDynamics,
     apf::AuxiliaryParticleFilter{<:AbstractParticleFilter},
     iter::Integer,
-    particle::Particle;
+    particle::AbstractParticle;
     kwargs...,
 )
     x_star = predictive_statistic(rng, apf.pp, dyn, iter, particle.state; kwargs...)
-    return Particle(x_star, particle.log_w, particle.ancestor)
+    return Particle(x_star, log_weight(particle), particle.ancestor)
 end
 
 function predictive_loglik(

GeneralisedFilters/src/algorithms/rbpf.jl

@@ -16,18 +16,20 @@ resampler(algo::RBPF) = resampler(algo.pf)
 function initialise_particle(
     rng::AbstractRNG, prior::HierarchicalPrior, algo::RBPF, ref_state; kwargs...
 )
+    N = num_particles(algo)
     x = sample_prior(rng, prior.outer_prior, algo.pf, ref_state; kwargs...)
     z = initialise(rng, prior.inner_prior, algo.af; new_outer=x, kwargs...)
     # TODO (RB):  determine the correct type for the log_w field or use a NoWeight type
-    return Particle(RBState(x, z), 0.0, 0)
+    # return Particle(RBState(x, z), -log(N), 0)
+    return Particle(RBState(x, z), 0)
 end
 
 function predict_particle(
     rng::AbstractRNG,
     dyn::HierarchicalDynamics,
     algo::RBPF,
     iter::Integer,
-    particle::RBParticle,
+    particle::AbstractParticle{<:RBState},
     observation,
     ref_state;
     kwargs...,
@@ -55,14 +57,16 @@ function predict_particle(
         kwargs...,
     )
 
-    return Particle(RBState(new_x, new_z), particle.log_w + logw_inc, particle.ancestor)
+    return Particle(
+        RBState(new_x, new_z), log_weight(particle) + logw_inc, particle.ancestor
+    )
 end
 
 function update_particle(
     obs::ObservationProcess,
     algo::RBPF,
     iter::Integer,
-    particle::RBParticle,
+    particle::AbstractParticle{<:RBState},
     observation;
     kwargs...,
 )
@@ -76,7 +80,9 @@ function update_particle(
         kwargs...,
     )
     return Particle(
-        RBState(particle.state.x, new_z), particle.log_w + log_increment, particle.ancestor
+        RBState(particle.state.x, new_z),
+        log_weight(particle) + log_increment,
+        particle.ancestor,
     )
 end
 
@@ -85,7 +91,7 @@ function predictive_state(
     dyn::HierarchicalDynamics,
     apf::AuxiliaryParticleFilter{<:RBPF},
     iter::Integer,
-    particle::RBParticle;
+    particle::AbstractParticle{<:RBState};
     kwargs...,
 )
     rbpf = apf.pf
@@ -110,7 +116,7 @@ function predictive_loglik(
     obs::ObservationProcess,
     algo::RBPF,
     iter::Integer,
-    p_star::RBParticle,
+    p_star::AbstractParticle{<:RBState},
     observation;
     kwargs...,
 )

GeneralisedFilters/src/containers.jl

@@ -2,18 +2,32 @@
 
 ## PARTICLES ###############################################################################
 
+abstract type AbstractParticle{T} end
+
 """
     Particle
 
 A container representing a single particle in a particle filter distribution, composed of a
 weighted sampled (stored as a log weight) and its ancestor index.
 """
-mutable struct Particle{ST,WT,AT<:Integer}
+mutable struct Particle{ST,WT,AT<:Integer} <: AbstractParticle{ST}
     state::ST
     log_w::WT
     ancestor::AT
 end
 
+# NOTE: this is only ever used for initializing a particle filter
+const UnweightedParticle{ST,AT} = Particle{ST,Nothing,AT}
+
+Particle(state, ancestor) = Particle(state, nothing, ancestor)
+Particle(particle::UnweightedParticle, ancestor) = Particle(particle.state, ancestor)
+function Particle(particle::Particle{<:Any,WT}, ancestor) where {WT<:Real}
+    return Particle(particle.state, zero(WT), ancestor)
+end
+
+log_weight(p::Particle{<:Any,<:Real}) = p.log_w
+log_weight(::UnweightedParticle) = false
+
 """
     RBState
 
@@ -30,8 +44,6 @@ mutable struct RBState{XT,ZT}
     z::ZT
 end
 
-const RBParticle{XT,ZT,WT} = Particle{RBState{XT,ZT},WT}
-
 """
     ParticleDistribution
 
@@ -44,7 +56,7 @@ their ancestories) into a distibution-like object.
   the unnormalized logsumexp of weights before update (for standard PF/guided filters)
   or a modified value that includes APF first-stage correction (for auxiliary PF).
 """
-mutable struct ParticleDistribution{WT,PT<:Particle{<:Any,WT},VT<:AbstractVector{PT}}
+mutable struct ParticleDistribution{WT,PT<:AbstractParticle,VT<:AbstractVector{PT}}
     particles::VT
     ll_baseline::WT
 end
@@ -59,10 +71,11 @@ Base.iterate(state::ParticleDistribution) = iterate(state.particles)
 # Not sure if this is kosher, since it doesn't follow the convention of Base.getindex
 Base.@propagate_inbounds Base.getindex(state::ParticleDistribution, i) = state.particles[i]
 
+log_weights(state::ParticleDistribution) = map(p -> log_weight(p), state.particles)
+get_weights(state::ParticleDistribution) = softmax(log_weights(state))
+
 # Helpers for StatsBase compatibility
-function StatsBase.weights(state::ParticleDistribution)
-    return Weights(softmax(map(p -> p.log_w, state.particles)))
-end
+StatsBase.weights(state::ParticleDistribution) = StatsBase.Weights(get_weights(state))
 
 """
     marginalise!(state::ParticleDistribution)
@@ -78,22 +91,21 @@ cases through a single-scalar caching mechanism. For standard PF, ll_baseline eq
 LSE before adding observation weights. For APF with resampling, it includes first-stage
 correction terms computed during the APF resampling step.
 """
-function marginalise!(state::ParticleDistribution)
+function marginalise!(state::ParticleDistribution, particles)
     # Compute logsumexp after adding observation likelihoods
-    LSE_after = logsumexp(map(p -> p.log_w, state.particles))
+    LSE_after = logsumexp(log_weight.(particles))
 
     # Compute log-likelihood increment: works for both PF and APF cases
     ll_increment = LSE_after - state.ll_baseline
 
     # Normalize weights
-    for p in state.particles
+    for p in particles
         p.log_w -= LSE_after
     end
 
     # Reset baseline for next iteration
-    state.ll_baseline = 0.0
-
-    return ll_increment
+    new_state = ParticleDistribution(particles, zero(ll_increment))
+    return new_state, ll_increment
 end
 
 ## GAUSSIAN STATES #########################################################################