Restore Type Stability (#122)

* first pass

* revert change to auxiliary filter

* added fix for APF

* correct likelihoods for APF

* Formatter

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

* added type stability with Float32s

* fixed typo

* added type stable `AuxiliaryResampler`

* introduce `UnweightedParticle` type

* minor adjustments

Just some fixes I noticed when going through my review

* merge conflict issue

* remove abstract particle

* fix example script

* revert `ScalMat`

* never resample uniform particles

* added typeless initial weights and normalizing constants

* formatter

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

* remove unused code

* fixed single precision static arrays

* added naive type stability tests

* Fix Aqua tests

* added `@test_opt` to unit tests

* cleaning before merge

---------

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

Author: 3 people

GeneralisedFilters/Project.toml

@@ -50,6 +50,7 @@ StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TestItemRunner = "f8b46487-2199-4994-9208-9a1283c18c0a"
 TestItems = "1c621080-faea-4a02-84b6-bbd5e436b8fe"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 
 [targets]
-test = ["Aqua", "PDMats", "StableRNGs", "Test", "TestItemRunner", "TestItems"]
+test = ["Aqua", "PDMats", "StableRNGs", "Test", "TestItemRunner", "TestItems", "JET"]

GeneralisedFilters/examples/trend-inflation/script.jl

@@ -14,8 +14,8 @@ using PDMats
 
 const GF = GeneralisedFilters
 
-# INFL_PATH = joinpath(@__DIR__, "..", "..", "..", "examples", "trend-inflation"); #hide
-INFL_PATH = joinpath(@__DIR__)
+INFL_PATH = joinpath(@__DIR__, "..", "..", "..", "examples", "trend-inflation"); #hide
+# INFL_PATH = joinpath(@__DIR__)
 include(joinpath(INFL_PATH, "utilities.jl")); #hide
 
 # ## Model Definition

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

@@ -46,14 +46,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, TypelessZero())
 end
 
 function predict(
@@ -66,20 +59,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(
@@ -93,10 +84,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
@@ -121,8 +111,7 @@ 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(
@@ -135,10 +124,10 @@ function predict_particle(
     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(
@@ -152,7 +141,7 @@ function update_particle(
     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(
@@ -252,10 +241,7 @@ 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
+    return new_x, TypelessZero()
 end
 
 # TODO: I feel like we shouldn't need to do this conversion. It should be handled by dispatch
@@ -353,7 +339,7 @@ function step(
     kwargs...,
 )
     # Compute lookahead weights approximating log p(y_{t+1} | x_{t}^(i))
-    log_ξs = map(state.particles) do particle
+    log_ηs = map(state.particles) do particle
         compute_logeta(
             rng,
             algo.weight_strategy,
@@ -366,40 +352,8 @@ function step(
         )
     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(

GeneralisedFilters/src/algorithms/rbpf.jl

@@ -18,16 +18,15 @@ function initialise_particle(
 )
     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), 0)
 end
 
 function predict_particle(
     rng::AbstractRNG,
     dyn::HierarchicalDynamics,
     algo::RBPF,
     iter::Integer,
-    particle::RBParticle,
+    particle::Particle{<:RBState},
     observation,
     ref_state;
     kwargs...,
@@ -55,14 +54,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::Particle{<:RBState},
     observation;
     kwargs...,
 )
@@ -76,7 +77,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
 

GeneralisedFilters/src/containers.jl

@@ -1,5 +1,68 @@
+using LogExpFunctions
+
 """Containers used for storing representations of the filtering distribution."""
 
+## TYPELESS INITIALIZERS ###################################################################
+
+"""
+    TypelessZero
+
+A lazy promotion for uninitialized particle weights whos type is not yet known at the first
+simulation of a particle filter.
+"""
+struct TypelessZero <: Number end
+
+Base.convert(::Type{T}, ::TypelessZero) where {T<:Number} = zero(T)
+Base.convert(::Type{TypelessZero}, ::TypelessZero) = TypelessZero()
+
+Base.:+(::TypelessZero, ::TypelessZero) = TypelessZero()
+
+Base.promote_rule(::Type{TypelessZero}, ::Type{T}) where {T<:Number} = T
+Base.promote_rule(::Type{TypelessZero}, ::Type{TypelessZero}) = TypelessZero
+
+Base.zero(::TypelessZero) = TypelessZero()
+Base.zero(::Type{TypelessZero}) = TypelessZero()
+
+Base.iszero(::TypelessZero) = true
+Base.isone(::TypelessZero) = false
+
+Base.show(io::IO, ::TypelessZero) = print(io, "TypelessZero()")
+
+"""
+    TypelessBaseline
+
+A lazy promotion for the computation of log-likelihood baslines given a collection of
+unweighted particles.
+"""
+struct TypelessBaseline <: Number
+    N::Int64
+end
+
+# Constructors for compatibility with Base.Number
+TypelessBaseline(x::TypelessBaseline) = x
+TypelessBaseline(x::Base.TwicePrecision) = TypelessBaseline(Int64(x))
+TypelessBaseline(x::AbstractChar) = TypelessBaseline(Int64(x))
+
+Base.convert(::Type{T}, b::TypelessBaseline) where {T<:Number} = T(log(b.N))
+Base.promote_rule(::Type{TypelessBaseline}, ::Type{T}) where {T<:Number} = T
+
+Base.iszero(::TypelessBaseline) = false
+Base.isone(::TypelessBaseline) = false
+
+function LogExpFunctions.logsumexp(weights::AbstractVector{TypelessZero})
+    return TypelessBaseline(length(weights))
+end
+
+function LogExpFunctions.softmax(x::AbstractVector{TypelessZero})
+    # TODO: horrible, but theoretically never used... except in the unit tests
+    return fill(1 / length(x), length(x))
+end
+
+Base.:+(::TypelessZero, b::TypelessBaseline) = b
+Base.:+(b::TypelessBaseline, ::TypelessZero) = b
+
+Base.show(io::IO, b::TypelessBaseline) = print(io, "Typeless(log($(b.N)))")
+
 ## PARTICLES ###############################################################################
 
 """
@@ -14,6 +77,17 @@ mutable struct Particle{ST,WT,AT<:Integer}
     ancestor::AT
 end
 
+# NOTE: this is only ever used for initializing a particle filter
+const UnweightedParticle{ST,AT} = Particle{ST,TypelessZero,AT}
+
+Particle(state, ancestor) = Particle(state, TypelessZero(), 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) = p.log_w
+
 """
     RBState
 
@@ -30,8 +104,6 @@ mutable struct RBState{XT,ZT}
     z::ZT
 end
 
-const RBParticle{XT,ZT,WT} = Particle{RBState{XT,ZT},WT}
-
 """
     ParticleDistribution
 
@@ -44,7 +116,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<:Particle,VT<:AbstractVector{PT}}
     particles::VT
     ll_baseline::WT
 end
@@ -59,10 +131,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 +151,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 #########################################################################